[flake8-pyi] Teach various rules that annotations might be stringized

Centralize the parsing of string annotations in a method on `Checker`
astral-sh · Sep 2, 2024 · 8c2653c · 8c2653c
1 parent 591a7a1
commit 8c2653c
Show file tree

Hide file tree

Showing 15 changed files with 191 additions and 51 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI032.py b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI032.py
@@ -3,8 +3,8 @@
 
 
 class Bad:
-    def __eq__(self, other: Any) -> bool: ...  # Y032
-    def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
+    def __eq__(self, other: Any) -> bool: ...  # PYI032
+    def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
 
 
 class Good:
@@ -22,3 +22,7 @@ class Unannotated:
   def __eq__(self) -> Any: ...
   def __ne__(self) -> bool: ...
 
+
+class BadStringized:
+   def __eq__(self, other: "Any") -> bool: ...  # PYI032
+   def __ne__(self, other: "Any") -> bool: ...  # PYI032
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI032.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI032.pyi
@@ -3,8 +3,8 @@ import typing
 
 
 class Bad:
-    def __eq__(self, other: Any) -> bool: ...  # Y032
-    def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
+    def __eq__(self, other: Any) -> bool: ...  # PYI032
+    def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
 
 
 class Good:
@@ -22,3 +22,6 @@ class Unannotated:
   def __eq__(self) -> Any: ...
   def __ne__(self) -> bool: ...
 
+class BadStringized:
+   def __eq__(self, other: "Any") -> bool: ...  # PYI032
+   def __ne__(self, other: "Any") -> bool: ...  # PYI032
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI034.py b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI034.py
@@ -317,3 +317,7 @@ def __ne__(self, other: Any) -> bool:
 
 def __imul__(self, other: Any) -> list[str]:
     ...
+
+class UsesStringizedAnnotations:
+    def __iadd__(self, other: "UsesStringizedAnnotations") -> "typing.Self":
+        return self
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI034.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI034.pyi
@@ -212,3 +212,6 @@ def __str__(self) -> str: ...
 def __eq__(self, other: Any) -> bool: ...
 def __ne__(self, other: Any) -> bool: ...
 def __imul__(self, other: Any) -> list[str]: ...
+
+class UsesStringizedAnnotations:
+    def __iadd__(self, other: "UsesStringizedAnnotations") -> "typing.Self": ...
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI050.py b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI050.py
@@ -30,3 +30,11 @@ def foo_no_return_pos_only(arg: int, /, arg2: NoReturn):
 
 def foo_never(arg: Never):
     ...
+
+
+def stringized(arg: "NoReturn"):
+    ...
+
+
+def stringized_good(arg: "Never"):
+    ...
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI050.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI050.pyi
@@ -21,3 +21,5 @@ def foo_int_kwargs_no_return(*args: NoReturn, **kwargs: int): ...  # Error: PYI0
 def foo_args_never(*args: Never): ...
 def foo_kwargs_never(**kwargs: Never): ...
 def foo_args_kwargs_never(*args: Never, **kwargs: Never): ...
+def stringized(arg: "NoReturn"): ...
+def stringized_good(arg: "Never"): ...
diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/any_eq_ne_annotation.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/any_eq_ne_annotation.rs
@@ -4,6 +4,7 @@ use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_text_size::Ranged;
 
+use super::helpers::match_maybe_stringized_annotation;
 use crate::checkers::ast::Checker;
 
 /// ## What it does
@@ -78,23 +79,27 @@ pub(crate) fn any_eq_ne_annotation(checker: &mut Checker, name: &str, parameters
         return;
     }
 
-    if semantic.match_typing_expr(annotation, "Any") {
-        let mut diagnostic = Diagnostic::new(
-            AnyEqNeAnnotation {
-                method_name: name.to_string(),
-            },
-            annotation.range(),
-        );
-        // Ex) `def __eq__(self, obj: Any): ...`
-        diagnostic.try_set_fix(|| {
-            let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
-                "object",
-                annotation.start(),
-                semantic,
-            )?;
-            let binding_edit = Edit::range_replacement(binding, annotation.range());
-            Ok(Fix::safe_edits(binding_edit, import_edit))
-        });
-        checker.diagnostics.push(diagnostic);
+    if !match_maybe_stringized_annotation(annotation, checker, |expr| {
+        semantic.match_typing_expr(expr, "Any")
+    }) {
+        return;
     }
+
+    let mut diagnostic = Diagnostic::new(
+        AnyEqNeAnnotation {
+            method_name: name.to_string(),
+        },
+        annotation.range(),
+    );
+    // Ex) `def __eq__(self, obj: Any): ...`
+    diagnostic.try_set_fix(|| {
+        let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
+            "object",
+            annotation.start(),
+            semantic,
+        )?;
+        let binding_edit = Edit::range_replacement(binding, annotation.range());
+        Ok(Fix::safe_edits(binding_edit, import_edit))
+    });
+    checker.diagnostics.push(diagnostic);
 }
diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/helpers.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/helpers.rs
@@ -0,0 +1,23 @@
+use ruff_python_ast as ast;
+
+use crate::checkers::ast::Checker;
+
+/// Apply a test to an annotation expression,
+/// abstracting over the fact that the annotation expression might be "stringized".
+///
+/// A stringized annotation is one enclosed in string quotes:
+/// `foo: "typing.Any"` means the same thing to a type checker as `foo: typing.Any`.
+pub(super) fn match_maybe_stringized_annotation(
+    expr: &ast::Expr,
+    checker: &Checker,
+    match_fn: impl FnOnce(&ast::Expr) -> bool,
+) -> bool {
+    if let ast::Expr::StringLiteral(string_annotation) = expr {
+        let Some(parsed_annotation) = checker.parse_type_annotation(string_annotation) else {
+            return false;
+        };
+        match_fn(parsed_annotation.expression())
+    } else {
+        match_fn(expr)
+    }
+}
diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/mod.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/mod.rs
@@ -56,6 +56,7 @@ mod ellipsis_in_non_empty_class_body;
 mod exit_annotations;
 mod future_annotations_in_stub;
 mod generic_not_last_base_class;
+mod helpers;
 mod iter_method_return_iterable;
 mod no_return_argument_annotation;
 mod non_empty_stub_body;

diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/no_return_argument_annotation.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/no_return_argument_annotation.rs
@@ -2,9 +2,10 @@ use std::fmt;
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::{AnyParameterRef, Parameters};
+use ruff_python_ast as ast;
 use ruff_text_size::Ranged;
 
+use super::helpers::match_maybe_stringized_annotation;
 use crate::checkers::ast::Checker;
 use crate::settings::types::PythonVersion::Py311;
 
@@ -54,14 +55,17 @@ impl Violation for NoReturnArgumentAnnotationInStub {
 }
 
 /// PYI050
-pub(crate) fn no_return_argument_annotation(checker: &mut Checker, parameters: &Parameters) {
+pub(crate) fn no_return_argument_annotation(checker: &mut Checker, parameters: &ast::Parameters) {
     // Ex) def func(arg: NoReturn): ...
     // Ex) def func(arg: NoReturn, /): ...
     // Ex) def func(*, arg: NoReturn): ...
     // Ex) def func(*args: NoReturn): ...
     // Ex) def func(**kwargs: NoReturn): ...
-    for annotation in parameters.iter().filter_map(AnyParameterRef::annotation) {
-        if checker.semantic().match_typing_expr(annotation, "NoReturn") {
+    for annotation in parameters
+        .iter()
+        .filter_map(ast::AnyParameterRef::annotation)
+    {
+        if is_no_return(annotation, checker) {
             checker.diagnostics.push(Diagnostic::new(
                 NoReturnArgumentAnnotationInStub {
                     module: if checker.settings.target_version >= Py311 {
@@ -76,6 +80,12 @@ pub(crate) fn no_return_argument_annotation(checker: &mut Checker, parameters: &
     }
 }
 
+fn is_no_return(expr: &ast::Expr, checker: &Checker) -> bool {
+    match_maybe_stringized_annotation(expr, checker, |expr| {
+        checker.semantic().match_typing_expr(expr, "NoReturn")
+    })
+}
+
 #[derive(Debug, PartialEq, Eq)]
 enum TypingModule {
     Typing,

diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/non_self_return_type.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/non_self_return_type.rs
@@ -8,6 +8,7 @@ use ruff_python_semantic::analyze;
 use ruff_python_semantic::analyze::visibility::{is_abstract, is_final, is_overload};
 use ruff_python_semantic::{ScopeKind, SemanticModel};
 
+use super::helpers::match_maybe_stringized_annotation;
 use crate::checkers::ast::Checker;
 
 /// ## What it does
@@ -149,7 +150,7 @@ pub(crate) fn non_self_return_type(
 
     // In-place methods that are expected to return `Self`.
     if is_inplace_bin_op(name) {
-        if !is_self(returns, semantic) {
+        if !is_self(returns, checker) {
             checker.diagnostics.push(Diagnostic::new(
                 NonSelfReturnType {
                     class_name: class_def.name.to_string(),
@@ -235,8 +236,10 @@ fn is_name(expr: &Expr, name: &str) -> bool {
 }
 
 /// Return `true` if the given expression resolves to `typing.Self`.
-fn is_self(expr: &Expr, semantic: &SemanticModel) -> bool {
-    semantic.match_typing_expr(expr, "Self")
+fn is_self(expr: &Expr, checker: &Checker) -> bool {
+    match_maybe_stringized_annotation(expr, checker, |expr| {
+        checker.semantic().match_typing_expr(expr, "Self")
+    })
 }
 
 /// Return `true` if the given class extends `collections.abc.Iterator`.

diff --git a/...c/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI032_PYI032.py.snap b/...c/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI032_PYI032.py.snap
@@ -4,39 +4,70 @@ source: crates/ruff_linter/src/rules/flake8_pyi/mod.rs
 PYI032.py:6:29: PYI032 [*] Prefer `object` to `Any` for the second parameter to `__eq__`
   |
 5 | class Bad:
-6 |     def __eq__(self, other: Any) -> bool: ...  # Y032
+6 |     def __eq__(self, other: Any) -> bool: ...  # PYI032
   |                             ^^^ PYI032
-7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
+7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
   |
   = help: Replace with `object`
 
 ℹ Safe fix
 3 3 | 
 4 4 | 
 5 5 | class Bad:
-6   |-    def __eq__(self, other: Any) -> bool: ...  # Y032
-  6 |+    def __eq__(self, other: object) -> bool: ...  # Y032
-7 7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
+6   |-    def __eq__(self, other: Any) -> bool: ...  # PYI032
+  6 |+    def __eq__(self, other: object) -> bool: ...  # PYI032
+7 7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
 8 8 | 
 9 9 | 
 
 PYI032.py:7:29: PYI032 [*] Prefer `object` to `Any` for the second parameter to `__ne__`
   |
 5 | class Bad:
-6 |     def __eq__(self, other: Any) -> bool: ...  # Y032
-7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
+6 |     def __eq__(self, other: Any) -> bool: ...  # PYI032
+7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
   |                             ^^^^^^^^^^ PYI032
   |
   = help: Replace with `object`
 
 ℹ Safe fix
 4 4 | 
 5 5 | class Bad:
-6 6 |     def __eq__(self, other: Any) -> bool: ...  # Y032
-7   |-    def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
-  7 |+    def __ne__(self, other: object) -> typing.Any: ...  # Y032
+6 6 |     def __eq__(self, other: Any) -> bool: ...  # PYI032
+7   |-    def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
+  7 |+    def __ne__(self, other: object) -> typing.Any: ...  # PYI032
 8 8 | 
 9 9 | 
 10 10 | class Good:
 
+PYI032.py:27:28: PYI032 [*] Prefer `object` to `Any` for the second parameter to `__eq__`
+   |
+26 | class BadStringized:
+27 |    def __eq__(self, other: "Any") -> bool: ...  # PYI032
+   |                            ^^^^^ PYI032
+28 |    def __ne__(self, other: "Any") -> bool: ...  # PYI032
+   |
+   = help: Replace with `object`
 
+ℹ Safe fix
+24 24 | 
+25 25 | 
+26 26 | class BadStringized:
+27    |-   def __eq__(self, other: "Any") -> bool: ...  # PYI032
+   27 |+   def __eq__(self, other: object) -> bool: ...  # PYI032
+28 28 |    def __ne__(self, other: "Any") -> bool: ...  # PYI032
+
+PYI032.py:28:28: PYI032 [*] Prefer `object` to `Any` for the second parameter to `__ne__`
+   |
+26 | class BadStringized:
+27 |    def __eq__(self, other: "Any") -> bool: ...  # PYI032
+28 |    def __ne__(self, other: "Any") -> bool: ...  # PYI032
+   |                            ^^^^^ PYI032
+   |
+   = help: Replace with `object`
+
+ℹ Safe fix
+25 25 | 
+26 26 | class BadStringized:
+27 27 |    def __eq__(self, other: "Any") -> bool: ...  # PYI032
+28    |-   def __ne__(self, other: "Any") -> bool: ...  # PYI032
+   28 |+   def __ne__(self, other: object) -> bool: ...  # PYI032
diff --git a/.../rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI032_PYI032.pyi.snap b/.../rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI032_PYI032.pyi.snap
@@ -4,39 +4,70 @@ source: crates/ruff_linter/src/rules/flake8_pyi/mod.rs
 PYI032.pyi:6:29: PYI032 [*] Prefer `object` to `Any` for the second parameter to `__eq__`
   |
 5 | class Bad:
-6 |     def __eq__(self, other: Any) -> bool: ...  # Y032
+6 |     def __eq__(self, other: Any) -> bool: ...  # PYI032
   |                             ^^^ PYI032
-7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
+7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
   |
   = help: Replace with `object`
 
 ℹ Safe fix
 3 3 | 
 4 4 | 
 5 5 | class Bad:
-6   |-    def __eq__(self, other: Any) -> bool: ...  # Y032
-  6 |+    def __eq__(self, other: object) -> bool: ...  # Y032
-7 7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
+6   |-    def __eq__(self, other: Any) -> bool: ...  # PYI032
+  6 |+    def __eq__(self, other: object) -> bool: ...  # PYI032
+7 7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
 8 8 | 
 9 9 | 
 
 PYI032.pyi:7:29: PYI032 [*] Prefer `object` to `Any` for the second parameter to `__ne__`
   |
 5 | class Bad:
-6 |     def __eq__(self, other: Any) -> bool: ...  # Y032
-7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
+6 |     def __eq__(self, other: Any) -> bool: ...  # PYI032
+7 |     def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
   |                             ^^^^^^^^^^ PYI032
   |
   = help: Replace with `object`
 
 ℹ Safe fix
 4 4 | 
 5 5 | class Bad:
-6 6 |     def __eq__(self, other: Any) -> bool: ...  # Y032
-7   |-    def __ne__(self, other: typing.Any) -> typing.Any: ...  # Y032
-  7 |+    def __ne__(self, other: object) -> typing.Any: ...  # Y032
+6 6 |     def __eq__(self, other: Any) -> bool: ...  # PYI032
+7   |-    def __ne__(self, other: typing.Any) -> typing.Any: ...  # PYI032
+  7 |+    def __ne__(self, other: object) -> typing.Any: ...  # PYI032
 8 8 | 
 9 9 | 
 10 10 | class Good:
 
+PYI032.pyi:26:28: PYI032 [*] Prefer `object` to `Any` for the second parameter to `__eq__`
+   |
+25 | class BadStringized:
+26 |    def __eq__(self, other: "Any") -> bool: ...  # PYI032
+   |                            ^^^^^ PYI032
+27 |    def __ne__(self, other: "Any") -> bool: ...  # PYI032
+   |
+   = help: Replace with `object`
 
+ℹ Safe fix
+23 23 |   def __ne__(self) -> bool: ...
+24 24 | 
+25 25 | class BadStringized:
+26    |-   def __eq__(self, other: "Any") -> bool: ...  # PYI032
+   26 |+   def __eq__(self, other: object) -> bool: ...  # PYI032
+27 27 |    def __ne__(self, other: "Any") -> bool: ...  # PYI032
+
+PYI032.pyi:27:28: PYI032 [*] Prefer `object` to `Any` for the second parameter to `__ne__`
+   |
+25 | class BadStringized:
+26 |    def __eq__(self, other: "Any") -> bool: ...  # PYI032
+27 |    def __ne__(self, other: "Any") -> bool: ...  # PYI032
+   |                            ^^^^^ PYI032
+   |
+   = help: Replace with `object`
+
+ℹ Safe fix
+24 24 | 
+25 25 | class BadStringized:
+26 26 |    def __eq__(self, other: "Any") -> bool: ...  # PYI032
+27    |-   def __ne__(self, other: "Any") -> bool: ...  # PYI032
+   27 |+   def __ne__(self, other: object) -> bool: ...  # PYI032