From bfad7cb1db9d64fc80b1820dc124c46a1b516014 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Tue, 21 Nov 2023 19:18:23 +0000
Subject: [PATCH] Respect dictionary unpacking in NamedTuple assignments

---
 .../test/fixtures/pyflakes/F821_23.py         |  7 +++
 crates/ruff_linter/src/checkers/ast/mod.rs    | 59 +++++++++++++------
 crates/ruff_linter/src/rules/pyflakes/mod.rs  |  1 +
 ...les__pyflakes__tests__F821_F821_23.py.snap |  4 ++
 4 files changed, 52 insertions(+), 19 deletions(-)
 create mode 100644 crates/ruff_linter/resources/test/fixtures/pyflakes/F821_23.py
 create mode 100644 crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F821_F821_23.py.snap

diff --git a/crates/ruff_linter/resources/test/fixtures/pyflakes/F821_23.py b/crates/ruff_linter/resources/test/fixtures/pyflakes/F821_23.py
new file mode 100644
index 0000000000000..456b9d1a7abd0
--- /dev/null
+++ 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
index 20b9ea2756b43..0a6ee16b15354 100644
--- 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
index 9f105649915cd..435755ce9fe42 100644
--- 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/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F821_F821_23.py.snap b/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F821_F821_23.py.snap
new file mode 100644
index 0000000000000..d0b409f39ee0b
--- /dev/null
+++ b/crates/ruff_linter/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
+---
+