[flake8-pyi] Fix PYI049 false negatives on call-based TypedDicts

astral-sh · Jan 17, 2024 · d2bfd3d · d2bfd3d
1 parent 368e279
commit d2bfd3d
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 13 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
@@ -2,7 +2,7 @@ use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::map_subscript;
 use ruff_python_ast::{self as ast, Expr, Stmt};
-use ruff_python_semantic::Scope;
+use ruff_python_semantic::{Scope, SemanticModel};
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -313,11 +313,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() {
@@ -327,23 +332,42 @@ 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 {
+        Stmt::ClassDef(class_def @ ast::StmtClassDef { name, .. }) => {
+            if class_def.bases().iter().any(is_typeddict) {
+                Some(name)
+            } else {
+                None
+            }
+        }
+        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})
+   |
+