Visit `TypeVar` and `NewType` name arguments (#4627)

This commit is contained in:
Charlie Marsh 2023-05-24 10:10:15 -04:00 committed by GitHub
parent a59d252246
commit 31bddef98f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 64 additions and 31 deletions

View File

@ -3685,32 +3685,42 @@ where
match callable { match callable {
Some(Callable::Bool) => { Some(Callable::Bool) => {
self.visit_expr(func); self.visit_expr(func);
if !args.is_empty() { let mut args = args.iter();
self.visit_boolean_test(&args[0]); if let Some(arg) = args.next() {
self.visit_boolean_test(arg);
} }
for expr in args.iter().skip(1) { for arg in args {
self.visit_expr(expr); self.visit_expr(arg);
} }
} }
Some(Callable::Cast) => { Some(Callable::Cast) => {
self.visit_expr(func); self.visit_expr(func);
if !args.is_empty() { let mut args = args.iter();
self.visit_type_definition(&args[0]); if let Some(arg) = args.next() {
self.visit_type_definition(arg);
} }
for expr in args.iter().skip(1) { for arg in args {
self.visit_expr(expr); self.visit_expr(arg);
} }
} }
Some(Callable::NewType) => { Some(Callable::NewType) => {
self.visit_expr(func); self.visit_expr(func);
for expr in args.iter().skip(1) { let mut args = args.iter();
self.visit_type_definition(expr); if let Some(arg) = args.next() {
self.visit_non_type_definition(arg);
}
for arg in args {
self.visit_type_definition(arg);
} }
} }
Some(Callable::TypeVar) => { Some(Callable::TypeVar) => {
self.visit_expr(func); self.visit_expr(func);
for expr in args.iter().skip(1) { let mut args = args.iter();
self.visit_type_definition(expr); if let Some(arg) = args.next() {
self.visit_non_type_definition(arg);
}
for arg in args {
self.visit_type_definition(arg);
} }
for keyword in keywords { for keyword in keywords {
let Keyword { let Keyword {
@ -3731,24 +3741,30 @@ where
self.visit_expr(func); self.visit_expr(func);
// Ex) NamedTuple("a", [("a", int)]) // Ex) NamedTuple("a", [("a", int)])
if args.len() > 1 { let mut args = args.iter();
match &args[1] { if let Some(arg) = args.next() {
Expr::List(ast::ExprList { elts, .. }) self.visit_non_type_definition(arg);
| Expr::Tuple(ast::ExprTuple { elts, .. }) => { }
for elt in elts { for arg in args {
match elt { if let Expr::List(ast::ExprList { elts, .. })
Expr::List(ast::ExprList { elts, .. }) | Expr::Tuple(ast::ExprTuple { elts, .. }) = arg
| Expr::Tuple(ast::ExprTuple { elts, .. }) => { {
if elts.len() == 2 { for elt in elts {
self.visit_non_type_definition(&elts[0]); match elt {
self.visit_type_definition(&elts[1]); Expr::List(ast::ExprList { elts, .. })
} | Expr::Tuple(ast::ExprTuple { elts, .. })
} if elts.len() == 2 =>
_ => {} {
self.visit_non_type_definition(&elts[0]);
self.visit_type_definition(&elts[1]);
}
_ => {
self.visit_non_type_definition(elt);
} }
} }
} }
_ => {} } else {
self.visit_non_type_definition(arg);
} }
} }
@ -3762,12 +3778,16 @@ where
self.visit_expr(func); self.visit_expr(func);
// Ex) TypedDict("a", {"a": int}) // Ex) TypedDict("a", {"a": int})
if args.len() > 1 { let mut args = args.iter();
if let Some(arg) = args.next() {
self.visit_non_type_definition(arg);
}
for arg in args {
if let Expr::Dict(ast::ExprDict { if let Expr::Dict(ast::ExprDict {
keys, keys,
values, values,
range: _, range: _,
}) = &args[1] }) = arg
{ {
for key in keys.iter().flatten() { for key in keys.iter().flatten() {
self.visit_non_type_definition(key); self.visit_non_type_definition(key);
@ -3775,6 +3795,8 @@ where
for value in values { for value in values {
self.visit_type_definition(value); self.visit_type_definition(value);
} }
} else {
self.visit_non_type_definition(arg);
} }
} }
@ -3787,11 +3809,12 @@ where
Some(Callable::MypyExtension) => { Some(Callable::MypyExtension) => {
self.visit_expr(func); self.visit_expr(func);
if let Some(arg) = args.first() { let mut args = args.iter();
if let Some(arg) = args.next() {
// Ex) DefaultNamedArg(bool | None, name="some_prop_name") // Ex) DefaultNamedArg(bool | None, name="some_prop_name")
self.visit_type_definition(arg); self.visit_type_definition(arg);
for arg in args.iter().skip(1) { for arg in args {
self.visit_non_type_definition(arg); self.visit_non_type_definition(arg);
} }
for keyword in keywords { for keyword in keywords {

View File

@ -3689,6 +3689,16 @@ mod tests {
"#, "#,
&[], &[],
); );
flakes(
r#"
from typing import NewType
def f():
name = "x"
NewType(name, int)
"#,
&[],
);
} }
#[test] #[test]