Respect dictionary unpacking in NamedTuple assignments (#8810)

Closes #8803.
astral-sh · Nov 21, 2023 · 5373759 · 5373759
1 parent d9151b1
commit 5373759
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 19 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/pyflakes/F821_23.py b/crates/ruff_linter/resources/test/fixtures/pyflakes/F821_23.py
@@ -0,0 +1,7 @@
+"""Test for type annotation parsing in `NamedTuple`."""
+
+import typing
+
+User = typing.NamedTuple('User', **{'name': str, 'password': bytes})
+User = typing.NamedTuple('User', name=str, password=bytes)
+User = typing.NamedTuple('User', [('name', str), ('password', bytes)])
diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs
@@ -1014,33 +1014,54 @@ where
                         if let Some(arg) = args.next() {
                             self.visit_non_type_definition(arg);
                         }
+
                         for arg in args {
-                            if let Expr::List(ast::ExprList { elts, .. })
-                            | Expr::Tuple(ast::ExprTuple { elts, .. }) = arg
-                            {
-                                for elt in elts {
-                                    match elt {
-                                        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);
+                            match arg {
+                                // Ex) NamedTuple("a", [("a", int)])
+                                Expr::List(ast::ExprList { elts, .. })
+                                | Expr::Tuple(ast::ExprTuple { elts, .. }) => {
+                                    for elt in elts {
+                                        match elt {
+                                            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);
+                                _ => self.visit_non_type_definition(arg),
                             }
                         }
 
-                        // Ex) NamedTuple("a", a=int)
                         for keyword in keywords {
-                            let Keyword { value, .. } = keyword;
-                            self.visit_type_definition(value);
+                            let Keyword { arg, value, .. } = keyword;
+                            match (arg.as_ref(), value) {
+                                // Ex) NamedTuple("a", **{"a": int})
+                                (None, Expr::Dict(ast::ExprDict { keys, values, .. })) => {
+                                    for (key, value) in keys.iter().zip(values) {
+                                        if let Some(key) = key.as_ref() {
+                                            self.visit_non_type_definition(key);
+                                            self.visit_type_definition(value);
+                                        } else {
+                                            self.visit_non_type_definition(value);
+                                        }
+                                    }
+                                }
+                                // Ex) NamedTuple("a", **obj)
+                                (None, _) => {
+                                    self.visit_non_type_definition(value);
+                                }
+                                // Ex) NamedTuple("a", a=int)
+                                _ => {
+                                    self.visit_type_definition(value);
+                                }
+                            }
                         }
                     }
                     Some(typing::Callable::TypedDict) => {

diff --git a/crates/ruff_linter/src/rules/pyflakes/mod.rs b/crates/ruff_linter/src/rules/pyflakes/mod.rs
@@ -140,6 +140,7 @@ mod tests {
     #[test_case(Rule::UndefinedName, Path::new("F821_20.py"))]
     #[test_case(Rule::UndefinedName, Path::new("F821_21.py"))]
     #[test_case(Rule::UndefinedName, Path::new("F821_22.ipynb"))]
+    #[test_case(Rule::UndefinedName, Path::new("F821_23.py"))]
     #[test_case(Rule::UndefinedExport, Path::new("F822_0.py"))]
     #[test_case(Rule::UndefinedExport, Path::new("F822_1.py"))]
     #[test_case(Rule::UndefinedExport, Path::new("F822_2.py"))]

diff --git a/...er/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F821_F821_23.py.snap b/...er/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F821_F821_23.py.snap
@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/rules/pyflakes/mod.rs
+---
+