[flake8-pyi] Fix incorrect behaviour of `custom-typevar-return-type…

…` preview-mode autofix if `typing` was already imported (`PYI019`) (astral-sh#15853)
tjkuson · Jan 31, 2025 · 0d191a1 · 0d191a1
1 parent b2cb757
commit 0d191a1
Show file tree

Hide file tree

Showing 10 changed files with 153 additions and 119 deletions.
diff --git a/...ources/test/fixtures/flake8_pyi/PYI019.py → ...rces/test/fixtures/flake8_pyi/PYI019_0.py b/...ources/test/fixtures/flake8_pyi/PYI019.py → ...rces/test/fixtures/flake8_pyi/PYI019_0.py
diff --git a/...urces/test/fixtures/flake8_pyi/PYI019.pyi → ...ces/test/fixtures/flake8_pyi/PYI019_0.pyi b/...urces/test/fixtures/flake8_pyi/PYI019.pyi → ...ces/test/fixtures/flake8_pyi/PYI019_0.pyi
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019_1.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019_1.pyi
@@ -0,0 +1,4 @@
+import typing
+
+class F:
+    def m[S](self: S) -> S: ...
diff --git a/crates/ruff_linter/src/rules/flake8_pyi/mod.rs b/crates/ruff_linter/src/rules/flake8_pyi/mod.rs
@@ -136,8 +136,9 @@ mod tests {
         Ok(())
     }
 
-    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019.py"))]
-    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019.pyi"))]
+    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019_0.py"))]
+    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019_0.pyi"))]
+    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019_1.pyi"))]
     fn custom_classmethod_rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
         let diagnostics = test_path(
@@ -154,20 +155,26 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn custom_classmethod_rules_preview() -> Result<()> {
+    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019_0.pyi"))]
+    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019_1.pyi"))]
+    fn custom_classmethod_rules_preview(rule_code: Rule, path: &Path) -> Result<()> {
+        let snapshot = format!(
+            "preview_{}_{}",
+            rule_code.noqa_code(),
+            path.to_string_lossy()
+        );
         let diagnostics = test_path(
-            Path::new("flake8_pyi/PYI019.pyi"),
+            Path::new("flake8_pyi").join(path).as_path(),
             &settings::LinterSettings {
                 pep8_naming: pep8_naming::settings::Settings {
                     classmethod_decorators: vec!["foo_classmethod".to_string()],
                     ..pep8_naming::settings::Settings::default()
                 },
                 preview: PreviewMode::Enabled,
-                ..settings::LinterSettings::for_rule(Rule::CustomTypeVarReturnType)
+                ..settings::LinterSettings::for_rule(rule_code)
             },
         )?;
-        assert_messages!(diagnostics);
+        assert_messages!(snapshot, diagnostics);
         Ok(())
     }
 

diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/custom_type_var_return_type.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/custom_type_var_return_type.rs
@@ -1,5 +1,6 @@
 use crate::checkers::ast::Checker;
-use crate::importer::ImportRequest;
+use crate::importer::{ImportRequest, ResolutionError};
+use crate::settings::types::PythonVersion;
 use itertools::Itertools;
 use ruff_diagnostics::{Applicability, Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
@@ -9,7 +10,7 @@ use ruff_python_ast::{
 use ruff_python_semantic::analyze::function_type::{self, FunctionType};
 use ruff_python_semantic::analyze::visibility::{is_abstract, is_overload};
 use ruff_python_semantic::SemanticModel;
-use ruff_text_size::{Ranged, TextRange};
+use ruff_text_size::{Ranged, TextRange, TextSize};
 
 /// ## What it does
 /// Checks for methods that define a custom `TypeVar` for their return type
@@ -270,12 +271,8 @@ fn add_diagnostic(checker: &mut Checker, function_def: &ast::StmtFunctionDef, re
         returns.range(),
     );
 
-    // See `replace_custom_typevar_with_self`'s doc comment
-    if in_stub {
-        if let Some(fix) = replace_custom_typevar_with_self(checker, function_def, returns) {
-            diagnostic.set_fix(fix);
-        }
-    }
+    diagnostic
+        .try_set_optional_fix(|| replace_custom_typevar_with_self(checker, function_def, returns));
 
     checker.diagnostics.push(diagnostic);
 }
@@ -288,66 +285,65 @@ fn add_diagnostic(checker: &mut Checker, function_def: &ast::StmtFunctionDef, re
 /// * Replace other uses of the original type variable elsewhere in the signature with `Self`
 /// * Remove that type variable from the PEP 695 type parameter list
 ///
-/// This fix cannot be suggested for non-stubs,
-/// as a non-stub fix would have to deal with references in body/at runtime as well,
-/// which is substantially harder and requires a type-aware backend.
-///
 /// The fourth step above has the same problem.
 /// This function thus only does replacements for the simplest of cases
 /// and will mark the fix as unsafe if an annotation cannot be handled.
 fn replace_custom_typevar_with_self(
     checker: &Checker,
     function_def: &ast::StmtFunctionDef,
     returns: &Expr,
-) -> Option<Fix> {
+) -> anyhow::Result<Option<Fix>> {
     if checker.settings.preview.is_disabled() {
-        return None;
+        return Ok(None);
+    }
+
+    // This fix cannot be suggested for non-stubs,
+    // as a non-stub fix would have to deal with references in body/at runtime as well,
+    // which is substantially harder and requires a type-aware backend.
+    if !checker.source_type.is_stub() {
+        return Ok(None);
     }
 
     // Non-`Name` return annotations are not currently autofixed
-    let typevar_name = &returns.as_name_expr()?.id;
+    let Expr::Name(typevar_name) = &returns else {
+        return Ok(None);
+    };
+
+    let typevar_name = &typevar_name.id;
+
+    let (import_edit, self_symbol_binding) = import_self(checker, returns.start())?;
 
     let mut all_edits = vec![
-        replace_return_annotation_with_self(returns),
+        import_edit,
+        replace_return_annotation_with_self(self_symbol_binding, returns),
         remove_first_parameter_annotation(&function_def.parameters),
     ];
 
-    let edit = import_self(checker, returns.range())?;
-    all_edits.push(edit);
-
-    if let Some(edit) =
-        remove_typevar_declaration(function_def.type_params.as_deref(), typevar_name)
-    {
-        all_edits.push(edit);
-    }
+    all_edits.extend(remove_typevar_declaration(
+        function_def.type_params.as_deref(),
+        typevar_name,
+    ));
 
-    let (mut edits, fix_applicability) =
+    let (edits, fix_applicability) =
         replace_typevar_usages_with_self(&function_def.parameters, typevar_name);
-    all_edits.append(&mut edits);
+
+    all_edits.extend(edits);
 
     let (first, rest) = (all_edits.swap_remove(0), all_edits);
 
-    Some(Fix::applicable_edits(first, rest, fix_applicability))
+    Ok(Some(Fix::applicable_edits(first, rest, fix_applicability)))
 }
 
-fn import_self(checker: &Checker, return_range: TextRange) -> Option<Edit> {
-    // From PYI034's fix
-    let target_version = checker.settings.target_version.as_tuple();
-    let source_module = if target_version >= (3, 11) {
+fn import_self(checker: &Checker, position: TextSize) -> Result<(Edit, String), ResolutionError> {
+    // See also PYI034's fix
+    let source_module = if checker.settings.target_version >= PythonVersion::Py311 {
         "typing"
     } else {
         "typing_extensions"
     };
-
     let (importer, semantic) = (checker.importer(), checker.semantic());
     let request = ImportRequest::import_from(source_module, "Self");
-
-    let position = return_range.start();
-    let (edit, ..) = importer
-        .get_or_import_symbol(&request, position, semantic)
-        .ok()?;
-
-    Some(edit)
+    importer.get_or_import_symbol(&request, position, semantic)
 }
 
 fn remove_first_parameter_annotation(parameters: &Parameters) -> Edit {
@@ -362,8 +358,8 @@ fn remove_first_parameter_annotation(parameters: &Parameters) -> Edit {
     Edit::deletion(name_end, annotation_end)
 }
 
-fn replace_return_annotation_with_self(returns: &Expr) -> Edit {
-    Edit::range_replacement("Self".to_string(), returns.range())
+fn replace_return_annotation_with_self(self_symbol_binding: String, returns: &Expr) -> Edit {
+    Edit::range_replacement(self_symbol_binding, returns.range())
 }
 
 fn replace_typevar_usages_with_self(

diff --git a/..._flake8_pyi__tests__PYI019_PYI019.py.snap → ...lake8_pyi__tests__PYI019_PYI019_0.py.snap b/..._flake8_pyi__tests__PYI019_PYI019.py.snap → ...lake8_pyi__tests__PYI019_PYI019_0.py.snap
@@ -1,56 +1,56 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_pyi/mod.rs
 ---
-PYI019.py:7:62: PYI019 Methods like `__new__` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:7:62: PYI019 Methods like `__new__` should return `Self` instead of a custom `TypeVar`
   |
 6 | class BadClass:
 7 |     def __new__(cls: type[_S], *args: str, **kwargs: int) -> _S: ...  # PYI019
   |                                                              ^^ PYI019
   |
 
-PYI019.py:10:54: PYI019 Methods like `bad_instance_method` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:10:54: PYI019 Methods like `bad_instance_method` should return `Self` instead of a custom `TypeVar`
    |
 10 |     def bad_instance_method(self: _S, arg: bytes) -> _S: ...  # PYI019
    |                                                      ^^ PYI019
    |
 
-PYI019.py:14:54: PYI019 Methods like `bad_class_method` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:14:54: PYI019 Methods like `bad_class_method` should return `Self` instead of a custom `TypeVar`
    |
 13 |     @classmethod
 14 |     def bad_class_method(cls: type[_S], arg: int) -> _S: ...  # PYI019
    |                                                      ^^ PYI019
    |
 
-PYI019.py:18:55: PYI019 Methods like `bad_posonly_class_method` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:18:55: PYI019 Methods like `bad_posonly_class_method` should return `Self` instead of a custom `TypeVar`
    |
 17 |     @classmethod
 18 |     def bad_posonly_class_method(cls: type[_S], /) -> _S: ...  # PYI019
    |                                                       ^^ PYI019
    |
 
-PYI019.py:39:63: PYI019 Methods like `__new__` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:39:63: PYI019 Methods like `__new__` should return `Self` instead of a custom `TypeVar`
    |
 37 | # Python > 3.12
 38 | class PEP695BadDunderNew[T]:
 39 |   def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ...  # PYI019
    |                                                               ^ PYI019
    |
 
-PYI019.py:42:46: PYI019 Methods like `generic_instance_method` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:42:46: PYI019 Methods like `generic_instance_method` should return `Self` instead of a custom `TypeVar`
    |
 42 |   def generic_instance_method[S](self: S) -> S: ...  # PYI019
    |                                              ^ PYI019
    |
 
-PYI019.py:54:32: PYI019 Methods like `foo` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:54:32: PYI019 Methods like `foo` should return `Self` instead of a custom `TypeVar`
    |
 52 |    # in the settings for this test:
 53 |    @foo_classmethod
 54 |    def foo[S](cls: type[S]) -> S: ...  # PYI019
    |                                ^ PYI019
    |
 
-PYI019.py:61:48: PYI019 Methods like `__new__` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:61:48: PYI019 Methods like `__new__` should return `Self` instead of a custom `TypeVar`
    |
 59 | # Only .pyi gets fixes, no fixes for .py
 60 | class PEP695Fix:
@@ -60,7 +60,7 @@ PYI019.py:61:48: PYI019 Methods like `__new__` should return `Self` instead of a
 63 |     def __init_subclass__[S](cls: type[S]) -> S: ...
    |
 
-PYI019.py:63:47: PYI019 Methods like `__init_subclass__` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:63:47: PYI019 Methods like `__init_subclass__` should return `Self` instead of a custom `TypeVar`
    |
 61 |     def __new__[S: PEP695Fix](cls: type[S]) -> S: ...
 62 |
@@ -70,7 +70,7 @@ PYI019.py:63:47: PYI019 Methods like `__init_subclass__` should return `Self` in
 65 |     def __neg__[S: PEP695Fix](self: S) -> S: ...
    |
 
-PYI019.py:65:43: PYI019 Methods like `__neg__` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:65:43: PYI019 Methods like `__neg__` should return `Self` instead of a custom `TypeVar`
    |
 63 |     def __init_subclass__[S](cls: type[S]) -> S: ...
 64 |
@@ -80,7 +80,7 @@ PYI019.py:65:43: PYI019 Methods like `__neg__` should return `Self` instead of a
 67 |     def __pos__[S](self: S) -> S: ...
    |
 
-PYI019.py:67:32: PYI019 Methods like `__pos__` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:67:32: PYI019 Methods like `__pos__` should return `Self` instead of a custom `TypeVar`
    |
 65 |     def __neg__[S: PEP695Fix](self: S) -> S: ...
 66 |
@@ -90,7 +90,7 @@ PYI019.py:67:32: PYI019 Methods like `__pos__` should return `Self` instead of a
 69 |     def __add__[S: PEP695Fix](self: S, other: S) -> S: ...
    |
 
-PYI019.py:69:53: PYI019 Methods like `__add__` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:69:53: PYI019 Methods like `__add__` should return `Self` instead of a custom `TypeVar`
    |
 67 |     def __pos__[S](self: S) -> S: ...
 68 |
@@ -100,7 +100,7 @@ PYI019.py:69:53: PYI019 Methods like `__add__` should return `Self` instead of a
 71 |     def __sub__[S](self: S, other: S) -> S: ...
    |
 
-PYI019.py:71:42: PYI019 Methods like `__sub__` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:71:42: PYI019 Methods like `__sub__` should return `Self` instead of a custom `TypeVar`
    |
 69 |     def __add__[S: PEP695Fix](self: S, other: S) -> S: ...
 70 |
@@ -110,7 +110,7 @@ PYI019.py:71:42: PYI019 Methods like `__sub__` should return `Self` instead of a
 73 |     @classmethod
    |
 
-PYI019.py:74:59: PYI019 Methods like `class_method_bound` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:74:59: PYI019 Methods like `class_method_bound` should return `Self` instead of a custom `TypeVar`
    |
 73 |     @classmethod
 74 |     def class_method_bound[S: PEP695Fix](cls: type[S]) -> S: ...
@@ -119,7 +119,7 @@ PYI019.py:74:59: PYI019 Methods like `class_method_bound` should return `Self` i
 76 |     @classmethod
    |
 
-PYI019.py:77:50: PYI019 Methods like `class_method_unbound` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:77:50: PYI019 Methods like `class_method_unbound` should return `Self` instead of a custom `TypeVar`
    |
 76 |     @classmethod
 77 |     def class_method_unbound[S](cls: type[S]) -> S: ...
@@ -128,7 +128,7 @@ PYI019.py:77:50: PYI019 Methods like `class_method_unbound` should return `Self`
 79 |     def instance_method_bound[S: PEP695Fix](self: S) -> S: ...
    |
 
-PYI019.py:79:57: PYI019 Methods like `instance_method_bound` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:79:57: PYI019 Methods like `instance_method_bound` should return `Self` instead of a custom `TypeVar`
    |
 77 |     def class_method_unbound[S](cls: type[S]) -> S: ...
 78 |
@@ -138,7 +138,7 @@ PYI019.py:79:57: PYI019 Methods like `instance_method_bound` should return `Self
 81 |     def instance_method_unbound[S](self: S) -> S: ...
    |
 
-PYI019.py:81:48: PYI019 Methods like `instance_method_unbound` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:81:48: PYI019 Methods like `instance_method_unbound` should return `Self` instead of a custom `TypeVar`
    |
 79 |     def instance_method_bound[S: PEP695Fix](self: S) -> S: ...
 80 |
@@ -148,7 +148,7 @@ PYI019.py:81:48: PYI019 Methods like `instance_method_unbound` should return `Se
 83 |     def instance_method_bound_with_another_parameter[S: PEP695Fix](self: S, other: S) -> S: ...
    |
 
-PYI019.py:83:90: PYI019 Methods like `instance_method_bound_with_another_parameter` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:83:90: PYI019 Methods like `instance_method_bound_with_another_parameter` should return `Self` instead of a custom `TypeVar`
    |
 81 |     def instance_method_unbound[S](self: S) -> S: ...
 82 |
@@ -158,7 +158,7 @@ PYI019.py:83:90: PYI019 Methods like `instance_method_bound_with_another_paramet
 85 |     def instance_method_unbound_with_another_parameter[S](self: S, other: S) -> S: ...
    |
 
-PYI019.py:85:81: PYI019 Methods like `instance_method_unbound_with_another_parameter` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:85:81: PYI019 Methods like `instance_method_unbound_with_another_parameter` should return `Self` instead of a custom `TypeVar`
    |
 83 |     def instance_method_bound_with_another_parameter[S: PEP695Fix](self: S, other: S) -> S: ...
 84 |
@@ -168,7 +168,7 @@ PYI019.py:85:81: PYI019 Methods like `instance_method_unbound_with_another_param
 87 |     def multiple_type_vars[S, *Ts, T](self: S, other: S, /, *args: *Ts, a: T, b: list[T]) -> S: ...
    |
 
-PYI019.py:87:94: PYI019 Methods like `multiple_type_vars` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:87:94: PYI019 Methods like `multiple_type_vars` should return `Self` instead of a custom `TypeVar`
    |
 85 |     def instance_method_unbound_with_another_parameter[S](self: S, other: S) -> S: ...
 86 |
@@ -178,23 +178,23 @@ PYI019.py:87:94: PYI019 Methods like `multiple_type_vars` should return `Self` i
 89 |     def mixing_old_and_new_style_type_vars[T](self: _S695, a: T, b: T) -> _S695: ...
    |
 
-PYI019.py:89:75: PYI019 Methods like `mixing_old_and_new_style_type_vars` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:89:75: PYI019 Methods like `mixing_old_and_new_style_type_vars` should return `Self` instead of a custom `TypeVar`
    |
 87 |     def multiple_type_vars[S, *Ts, T](self: S, other: S, /, *args: *Ts, a: T, b: list[T]) -> S: ...
 88 |
 89 |     def mixing_old_and_new_style_type_vars[T](self: _S695, a: T, b: T) -> _S695: ...
    |                                                                           ^^^^^ PYI019
    |
 
-PYI019.py:102:40: PYI019 Methods like `m` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:102:40: PYI019 Methods like `m` should return `Self` instead of a custom `TypeVar`
     |
 100 | class UsesFullyQualifiedType:
 101 |     @classmethod
 102 |     def m[S](cls: builtins.type[S]) -> S: ...  # PYI019
     |                                        ^ PYI019
     |
 
-PYI019.py:114:31: PYI019 Methods like `m` should return `Self` instead of a custom `TypeVar`
+PYI019_0.py:114:31: PYI019 Methods like `m` should return `Self` instead of a custom `TypeVar`
     |
 112 | class SubscriptReturnType:
 113 |     @classmethod