Treat TypeAlias values as annotations (#377)

This commit is contained in:
Charlie Marsh 2022-10-09 17:37:19 -04:00 committed by GitHub
parent aafe7c0c39
commit 30b1b1e15a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 21 additions and 1 deletions

View file

@ -67,7 +67,22 @@ c = cast("Vegetable", b)
Field = lambda default=MISSING: field(default=default) Field = lambda default=MISSING: field(default=default)
# Test: access a sub-importation via an alias.
import pyarrow as pa import pyarrow as pa
import pyarrow.csv import pyarrow.csv
print(pa.csv.read_csv("test.csv")) print(pa.csv.read_csv("test.csv"))
# Test: referencing an import via TypeAlias.
import sys
import numpy as np
if sys.version_info >= (3, 10):
from typing import TypeAlias
else:
from typing_extensions import TypeAlias
CustomInt: TypeAlias = "np.int8 | np.int16"

View file

@ -1,3 +1,4 @@
use crate::ast::helpers::match_name_or_attr;
use rustpython_parser::ast::{ use rustpython_parser::ast::{
Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler, Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler,
ExcepthandlerKind, Expr, ExprContext, ExprKind, Keyword, MatchCase, Operator, Pattern, ExcepthandlerKind, Expr, ExprContext, ExprKind, Keyword, MatchCase, Operator, Pattern,
@ -148,8 +149,12 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
} => { } => {
visitor.visit_annotation(annotation); visitor.visit_annotation(annotation);
if let Some(expr) = value { if let Some(expr) = value {
if match_name_or_attr(annotation, "TypeAlias") {
visitor.visit_annotation(expr);
} else {
visitor.visit_expr(expr); visitor.visit_expr(expr);
} }
}
visitor.visit_expr(target); visitor.visit_expr(target);
} }
StmtKind::For { StmtKind::For {