[flake8-pyi] Fix PYI049 false negatives on call-based TypedDicts (#9567)

## Summary Fixes another of the bullet points from #8771 ## Test Plan `cargo test` / `cargo insta review`
astral-sh · Jan 18, 2024 · b3a6f0c · b3a6f0c
1 parent 7be7066
commit b3a6f0c
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 12 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI049.py b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI049.py
@@ -16,3 +16,8 @@ class _UsedTypedDict(TypedDict):
 
 class _CustomClass(_UsedTypedDict):
     bar: list[int]
+
+_UnusedTypedDict3 = TypedDict("_UnusedTypedDict3", {"foo": int})
+_UsedTypedDict3 = TypedDict("_UsedTypedDict3", {"bar": bytes})
+
+def uses_UsedTypedDict3(arg: _UsedTypedDict3) -> None: ...
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI049.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI049.pyi
@@ -30,3 +30,8 @@ else:
 
 class _CustomClass2(_UsedTypedDict2):
    bar: list[int]
+
+_UnusedTypedDict3 = TypedDict("_UnusedTypedDict3", {"foo": int})
+_UsedTypedDict3 = TypedDict("_UsedTypedDict3", {"bar": bytes})
+
+def uses_UsedTypedDict3(arg: _UsedTypedDict3) -> None: ...
diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/unused_private_type_definition.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/unused_private_type_definition.rs
@@ -323,11 +323,16 @@ pub(crate) fn unused_private_typed_dict(
     scope: &Scope,
     diagnostics: &mut Vec<Diagnostic>,
 ) {
+    let semantic = checker.semantic();
+
     for binding in scope
         .binding_ids()
-        .map(|binding_id| checker.semantic().binding(binding_id))
+        .map(|binding_id| semantic.binding(binding_id))
     {
-        if !(binding.kind.is_class_definition() && binding.is_private_declaration()) {
+        if !binding.is_private_declaration() {
+            continue;
+        }
+        if !(binding.kind.is_class_definition() || binding.kind.is_assignment()) {
             continue;
         }
         if binding.is_used() {
@@ -337,23 +342,64 @@ pub(crate) fn unused_private_typed_dict(
         let Some(source) = binding.source else {
             continue;
         };
-        let Stmt::ClassDef(class_def) = checker.semantic().statement(source) else {
-            continue;
-        };
 
-        if !class_def
-            .bases()
-            .iter()
-            .any(|base| checker.semantic().match_typing_expr(base, "TypedDict"))
-        {
+        let Some(class_name) = extract_typeddict_name(semantic.statement(source), semantic) else {
             continue;
-        }
+        };
 
         diagnostics.push(Diagnostic::new(
             UnusedPrivateTypedDict {
-                name: class_def.name.to_string(),
+                name: class_name.to_string(),
             },
             binding.range(),
         ));
     }
 }
+
+fn extract_typeddict_name<'a>(stmt: &'a Stmt, semantic: &SemanticModel) -> Option<&'a str> {
+    let is_typeddict = |expr: &ast::Expr| semantic.match_typing_expr(expr, "TypedDict");
+    match stmt {
+        // E.g. return `Some("Foo")` for the first one of these classes,
+        // and `Some("Bar")` for the second:
+        //
+        // ```python
+        // import typing
+        // from typing import TypedDict
+        //
+        // class Foo(TypedDict):
+        //     x: int
+        //
+        // T = typing.TypeVar("T")
+        //
+        // class Bar(typing.TypedDict, typing.Generic[T]):
+        //     y: T
+        // ```
+        Stmt::ClassDef(class_def @ ast::StmtClassDef { name, .. }) => {
+            if class_def.bases().iter().any(is_typeddict) {
+                Some(name)
+            } else {
+                None
+            }
+        }
+        // E.g. return `Some("Baz")` for this assignment,
+        // which is an accepted alternative way of creating a TypedDict type:
+        //
+        // ```python
+        // import typing
+        // Baz = typing.TypedDict("Baz", {"z": bytes})
+        // ```
+        Stmt::Assign(ast::StmtAssign { targets, value, .. }) => {
+            let [target] = targets.as_slice() else {
+                return None;
+            };
+            let ast::ExprName { id, .. } = target.as_name_expr()?;
+            let ast::ExprCall { func, .. } = value.as_call_expr()?;
+            if is_typeddict(func) {
+                Some(id)
+            } else {
+                None
+            }
+        }
+        _ => None,
+    }
+}
diff --git a/...c/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI049_PYI049.py.snap b/...c/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI049_PYI049.py.snap
@@ -15,4 +15,13 @@ PYI049.py:9:7: PYI049 Private TypedDict `_UnusedTypedDict2` is never used
 10 |     bar: int
    |
 
+PYI049.py:20:1: PYI049 Private TypedDict `_UnusedTypedDict3` is never used
+   |
+18 |     bar: list[int]
+19 | 
+20 | _UnusedTypedDict3 = TypedDict("_UnusedTypedDict3", {"foo": int})
+   | ^^^^^^^^^^^^^^^^^ PYI049
+21 | _UsedTypedDict3 = TypedDict("_UsedTypedDict3", {"bar": bytes})
+   |
+
 
diff --git a/.../rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI049_PYI049.pyi.snap b/.../rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI049_PYI049.pyi.snap
@@ -15,4 +15,13 @@ PYI049.pyi:10:7: PYI049 Private TypedDict `_UnusedTypedDict2` is never used
 11 |     bar: int
    |
 
+PYI049.pyi:34:1: PYI049 Private TypedDict `_UnusedTypedDict3` is never used
+   |
+32 |    bar: list[int]
+33 | 
+34 | _UnusedTypedDict3 = TypedDict("_UnusedTypedDict3", {"foo": int})
+   | ^^^^^^^^^^^^^^^^^ PYI049
+35 | _UsedTypedDict3 = TypedDict("_UsedTypedDict3", {"bar": bytes})
+   |
+