[flake8-pyi] Make PYI019 autofixable for .py files in preview m…

…ode as well as stubs

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019_0.py

@@ -1,4 +1,4 @@
-from typing import TypeVar, Self, Type
+from typing import TypeVar, Self, Type, cast
 
 _S = TypeVar("_S", bound=BadClass)
 _S2 = TypeVar("_S2", BadClass, GoodClass)
@@ -56,7 +56,7 @@ def foo[S](cls: type[S]) -> S: ...  # PYI019
 
 _S695 = TypeVar("_S695", bound="PEP695Fix")
 
-# Only .pyi gets fixes, no fixes for .py
+
 class PEP695Fix:
     def __new__[S: PEP695Fix](cls: type[S]) -> S: ...
 
@@ -135,3 +135,18 @@ class NoReturnAnnotations:
     def m[S](self: S, other: S): ...
     @classmethod
     def n[S](cls: type[S], other: S): ...
+
+class MethodsWithBody:
+    def m[S](self: S, other: S) -> S:
+        x: S = other
+        return x
+
+    @classmethod
+    def n[S](cls: type[S], other: S) -> S:
+        x: type[S] = type(other)
+        return x()
+
+class StringizedReferencesAreTooComplicated:
+    def m[S](self: S) -> S:
+        x = cast("S", self)
+        return x

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019_0.pyi

@@ -56,7 +56,7 @@ class CustomClassMethod:
 
 _S695 = TypeVar("_S695", bound="PEP695Fix")
 
-# Only .pyi gets fixes, no fixes for .py
+
 class PEP695Fix:
     def __new__[S: PEP695Fix](cls: type[S]) -> S: ...
 

crates/ruff_linter/src/rules/flake8_pyi/rules/custom_type_var_for_self.rs

@@ -13,8 +13,8 @@ use crate::importer::{ImportRequest, ResolutionError};
 use crate::settings::types::{PreviewMode, PythonVersion};
 
 /// ## What it does
-/// Checks for methods that use custom `TypeVar`s in their annotations
-/// when they could use `Self` instead.
+/// Checks for methods that use custom [`TypeVar`s][typing_TypeVar] in their
+/// annotations when they could use [`Self`][Self] instead.
 ///
 /// ## Why is this bad?
 /// While the semantics are often identical, using `Self` is more intuitive
@@ -48,10 +48,11 @@ use crate::settings::types::{PreviewMode, PythonVersion};
 ///     def bar(cls, arg: int) -> Self: ...
 /// ```
 ///
-/// ## Fix safety
-/// The fix is only available in stub files.
-/// It will try to remove all usages and declarations of the custom type variable.
-/// Pre-[PEP-695]-style declarations will not be removed.
+/// ## Fix behaviour and safety
+/// The fix removes all usages and declarations of the custom type variable.
+/// [PEP-695]-style `TypeVar` declarations are also removed from the [type parameter list];
+/// however, old-style `TypeVar`s do not have their declarations removed. See
+/// [`unused-private-type-var`][PYI018] for a rule to clean up unused private type variables.
 ///
 /// If there are any comments within the fix ranges, it will be marked as unsafe.
 /// Otherwise, it will be marked as safe.
@@ -70,6 +71,10 @@ use crate::settings::types::{PreviewMode, PythonVersion};
 ///
 /// [PEP 673]: https://peps.python.org/pep-0673/#motivation
 /// [PEP 695]: https://peps.python.org/pep-0695/
+/// [PYI018]: https://docs.astral.sh/ruff/rules/unused-private-type-var/
+/// [type parameter list]: https://docs.python.org/3/reference/compound_stmts.html#type-params
+/// [Self]: https://docs.python.org/3/library/typing.html#typing.Self
+/// [typing_TypeVar]: https://docs.python.org/3/library/typing.html#typing.TypeVar
 #[derive(ViolationMetadata)]
 pub(crate) struct CustomTypeVarForSelf {
     typevar_name: String,
@@ -195,7 +200,6 @@ pub(crate) fn custom_type_var_instead_of_self(
             &custom_typevar,
             self_or_cls_parameter,
             self_or_cls_annotation,
-            function_header_end,
         )
     });
 
@@ -460,19 +464,11 @@ fn replace_custom_typevar_with_self(
     custom_typevar: &TypeVar,
     self_or_cls_parameter: &ast::ParameterWithDefault,
     self_or_cls_annotation: &ast::Expr,
-    function_header_end: TextSize,
 ) -> anyhow::Result<Option<Fix>> {
     if checker.settings.preview.is_disabled() {
         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);
-    }
-
     let (import_edit, self_symbol_binding) = import_self(checker, function_def.start())?;
 
     let mut other_edits = vec![Edit::deletion(
@@ -489,16 +485,16 @@ fn replace_custom_typevar_with_self(
         other_edits.push(deletion_edit);
     }
 
-    let replace_references_range =
-        TextRange::new(self_or_cls_annotation.end(), function_header_end);
+    let replace_references_range = TextRange::new(self_or_cls_annotation.end(), function_def.end());
 
-    other_edits.extend(replace_typevar_usages_with_self(
+    replace_typevar_usages_with_self(
         custom_typevar,
         self_or_cls_annotation.range(),
         &self_symbol_binding,
         replace_references_range,
         checker.semantic(),
-    ));
+        &mut other_edits,
+    )?;
 
     let comment_ranges = checker.comment_ranges();
 
@@ -540,17 +536,27 @@ fn replace_typevar_usages_with_self<'a>(
     self_symbol_binding: &'a str,
     editable_range: TextRange,
     semantic: &'a SemanticModel<'a>,
-) -> impl Iterator<Item = Edit> + 'a {
-    typevar
-        .references(semantic)
-        .map(Ranged::range)
-        .filter(move |reference_range| editable_range.contains_range(*reference_range))
-        .filter(move |reference_range| {
-            !self_or_cls_annotation_range.contains_range(*reference_range)
-        })
-        .map(|reference_range| {
-            Edit::range_replacement(self_symbol_binding.to_string(), reference_range)
-        })
+    edits: &mut Vec<Edit>,
+) -> anyhow::Result<()> {
+    for reference in typevar.references(semantic) {
+        if reference.in_string_type_definition() {
+            bail!(
+                "Cannot apply autofix where some references to the TypeVar are in string type definitions"
+            );
+        }
+        let reference_range = reference.range();
+        if !editable_range.contains_range(reference_range) {
+            continue;
+        }
+        if self_or_cls_annotation_range.contains_range(reference_range) {
+            continue;
+        }
+        edits.push(Edit::range_replacement(
+            self_symbol_binding.to_string(),
+            reference_range,
+        ));
+    }
+    Ok(())
 }
 
 fn remove_pep695_typevar_declaration(

crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI019_PYI019_0.py.snap

@@ -59,7 +59,6 @@ PYI019_0.py:54:32: PYI019 Use `Self` instead of custom TypeVar `S`
 
 PYI019_0.py:61:48: PYI019 Use `Self` instead of custom TypeVar `S`
    |
-59 | # Only .pyi gets fixes, no fixes for .py
 60 | class PEP695Fix:
 61 |     def __new__[S: PEP695Fix](cls: type[S]) -> S: ...
    |                                                ^ PYI019
@@ -235,3 +234,33 @@ PYI019_0.py:131:40: PYI019 Use `Self` instead of custom TypeVar `_NotATypeVar`
     |                                        ^^^^^^^^^^^^ PYI019
     |
     = help: Replace TypeVar `_NotATypeVar` with `Self`
+
+PYI019_0.py:140:36: PYI019 Use `Self` instead of custom TypeVar `S`
+    |
+139 | class MethodsWithBody:
+140 |     def m[S](self: S, other: S) -> S:
+    |                                    ^ PYI019
+141 |         x: S = other
+142 |         return x
+    |
+    = help: Replace TypeVar `S` with `Self`
+
+PYI019_0.py:145:41: PYI019 Use `Self` instead of custom TypeVar `S`
+    |
+144 |     @classmethod
+145 |     def n[S](cls: type[S], other: S) -> S:
+    |                                         ^ PYI019
+146 |         x: type[S] = type(other)
+147 |         return x()
+    |
+    = help: Replace TypeVar `S` with `Self`
+
+PYI019_0.py:150:26: PYI019 Use `Self` instead of custom TypeVar `S`
+    |
+149 | class StringizedReferencesAreTooComplicated:
+150 |     def m[S](self: S) -> S:
+    |                          ^ PYI019
+151 |         x = cast("S", self)
+152 |         return x
+    |
+    = help: Replace TypeVar `S` with `Self`

crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI019_PYI019_0.pyi.snap

@@ -59,7 +59,6 @@ PYI019_0.pyi:54:32: PYI019 Use `Self` instead of custom TypeVar `S`
 
 PYI019_0.pyi:61:48: PYI019 Use `Self` instead of custom TypeVar `S`
    |
-59 | # Only .pyi gets fixes, no fixes for .py
 60 | class PEP695Fix:
 61 |     def __new__[S: PEP695Fix](cls: type[S]) -> S: ...
    |                                                ^ PYI019