[flake8-pyi] Add "replace with Self" fix (PYI019)

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

@@ -52,3 +52,34 @@ class CustomClassMethod:
    # in the settings for this test:
    @foo_classmethod
    def foo[S](cls: type[S]) -> S: ...  # PYI019
+
+
+# No fixes for .py
+class PEP695Fix:
+    def __new__[S: PEP695Fix](cls: type[S]) -> S: ...
+
+    def __init_subclass__[S](cls: type[S]) -> S: ...
+
+    def __neg__[S: PEP695Fix](self: S) -> S: ...
+
+    def __pos__[S](self: S) -> S: ...
+
+    def __add__[S: PEP695Fix](self: S, other: S) -> S: ...
+
+    def __sub__[S](self: S, other: S) -> S: ...
+
+    @classmethod
+    def class_method_bound[S: PEP695Fix](cls: type[S]) -> S: ...
+
+    @classmethod
+    def class_method_unbound[S](cls: type[S]) -> S: ...
+
+    def instance_method_bound[S: PEP695Fix](self: S) -> S: ...
+
+    def instance_method_unbound[S](self: S) -> S: ...
+
+    def instance_method_bound_with_another_parameter[S: PEP695Fix](self: S, other: S) -> S: ...
+
+    def instance_method_unbound_with_another_parameter[S](self: S, other: S) -> S: ...
+
+    def multiple_type_vars[S, *Ts, T](self: S, other: S, /, *args: *Ts, a: T, b: T) -> S: ...

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

@@ -52,3 +52,34 @@ class CustomClassMethod:
    # in the settings for this test:
    @foo_classmethod
    def foo[S](cls: type[S]) -> S: ...  # PYI019
+
+
+# Only .pyi gets fixes
+class PEP695Fix:
+    def __new__[S: PEP695Fix](cls: type[S]) -> S: ...
+
+    def __init_subclass__[S](cls: type[S]) -> S: ...
+
+    def __neg__[S: PEP695Fix](self: S) -> S: ...
+
+    def __pos__[S](self: S) -> S: ...
+
+    def __add__[S: PEP695Fix](self: S, other: S) -> S: ...
+
+    def __sub__[S](self: S, other: S) -> S: ...
+
+    @classmethod
+    def class_method_bound[S: PEP695Fix](cls: type[S]) -> S: ...
+
+    @classmethod
+    def class_method_unbound[S](cls: type[S]) -> S: ...
+
+    def instance_method_bound[S: PEP695Fix](self: S) -> S: ...
+
+    def instance_method_unbound[S](self: S) -> S: ...
+
+    def instance_method_bound_with_another_parameter[S: PEP695Fix](self: S, other: S) -> S: ...
+
+    def instance_method_unbound_with_another_parameter[S](self: S, other: S) -> S: ...
+
+    def multiple_type_vars[S, *Ts, T](self: S, other: S, /, *args: *Ts, a: T, b: T) -> S: ...

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -164,9 +164,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     checker,
                     name,
                     decorator_list,
-                    returns.as_ref().map(AsRef::as_ref),
-                    parameters,
                     type_params.as_deref(),
+                    parameters,
+                    returns.as_ref().map(AsRef::as_ref),
                 );
             }
             if checker.source_type.is_stub() {

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

@@ -1,26 +1,27 @@
-use ruff_diagnostics::{Diagnostic, Violation};
+use crate::checkers::ast::Checker;
+use crate::importer::ImportRequest;
+use itertools::Itertools;
+use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast as ast;
 use ruff_python_ast::helpers::map_subscript;
 use ruff_python_ast::{Decorator, Expr, Parameters, TypeParam, TypeParams};
 use ruff_python_semantic::analyze::function_type::{self, FunctionType};
 use ruff_python_semantic::analyze::visibility::{is_abstract, is_overload};
-use ruff_text_size::Ranged;
-
-use crate::checkers::ast::Checker;
+use ruff_text_size::{Ranged, TextRange};
 
 /// ## What it does
 /// Checks for methods that define a custom `TypeVar` for their return type
-/// annotation instead of using `typing_extensions.Self`.
+/// annotation instead of using `Self`.
 ///
 /// ## Why is this bad?
-/// While the semantics are often identical, using `typing_extensions.Self` is
-/// more intuitive and succinct (per [PEP 673]) than a custom `TypeVar`. For
-/// example, the use of `Self` will typically allow for the omission of type
-/// parameters on the `self` and `cls` arguments.
+/// While the semantics are often identical, using `Self` is more intuitive
+/// and succinct (per [PEP 673]) than a custom `TypeVar`. For example, the
+/// use of `Self` will typically allow for the omission of type parameters
+/// on the `self` and `cls` arguments.
 ///
-/// This check currently applies to instance methods that return `self`, class
-/// methods that return an instance of `cls`, and `__new__` methods.
+/// This check currently applies to instance methods that return `self`,
+/// class methods that return an instance of `cls`, and `__new__` methods.
 ///
 /// ## Example
 ///
@@ -48,15 +49,25 @@ use crate::checkers::ast::Checker;
 #[violation]
 pub struct CustomTypeVarReturnType {
     method_name: String,
+    in_stub: bool,
 }
 
 impl Violation for CustomTypeVarReturnType {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
-        let CustomTypeVarReturnType { method_name } = self;
-        format!(
-            "Methods like `{method_name}` should return `typing.Self` instead of a custom `TypeVar`"
-        )
+        let method_name = &self.method_name;
+        format!("Methods like `{method_name}` should return `Self` instead of a custom `TypeVar`")
+    }
+
+    fn fix_title(&self) -> Option<String> {
+        // See `replace_custom_typevar_with_self`'s doc comment
+        if self.in_stub {
+            Some("Replace with `Self`".to_string())
+        } else {
+            None
+        }
     }
 }
 
@@ -65,20 +76,20 @@ pub(crate) fn custom_type_var_return_type(
     checker: &mut Checker,
     name: &str,
     decorator_list: &[Decorator],
-    returns: Option<&Expr>,
-    args: &Parameters,
     type_params: Option<&TypeParams>,
+    parameters: &Parameters,
+    returns: Option<&Expr>,
 ) {
     // Given, e.g., `def foo(self: _S, arg: bytes) -> _T`, extract `_T`.
     let Some(returns) = returns else {
         return;
     };
 
     // Given, e.g., `def foo(self: _S, arg: bytes)`, extract `_S`.
-    let Some(self_or_cls_annotation) = args
+    let Some(self_or_cls_annotation) = parameters
         .posonlyargs
         .iter()
-        .chain(&args.args)
+        .chain(&parameters.args)
         .next()
         .and_then(|parameter_with_default| parameter_with_default.parameter.annotation.as_ref())
     else {
@@ -115,12 +126,7 @@ pub(crate) fn custom_type_var_return_type(
     };
 
     if method.uses_custom_var() {
-        checker.diagnostics.push(Diagnostic::new(
-            CustomTypeVarReturnType {
-                method_name: name.to_string(),
-            },
-            returns.range(),
-        ));
+        add_diagnostic(checker, name.to_string(), type_params, parameters, returns);
     }
 }
 
@@ -147,8 +153,8 @@ struct ClassMethod<'a> {
 }
 
 impl<'a> ClassMethod<'a> {
-    /// Returns `true` if the class method is annotated with a custom `TypeVar` that is likely
-    /// private.
+    /// Returns `true` if the class method is annotated with
+    /// a custom `TypeVar` that is likely private.
     fn uses_custom_var(&self) -> bool {
         let Expr::Subscript(ast::ExprSubscript { slice, value, .. }) = self.cls_annotation else {
             return false;
@@ -188,8 +194,8 @@ struct InstanceMethod<'a> {
 }
 
 impl<'a> InstanceMethod<'a> {
-    /// Returns `true` if the instance method is annotated with a custom `TypeVar` that is likely
-    /// private.
+    /// Returns `true` if the instance method is annotated with
+    /// a custom `TypeVar` that is likely private.
     fn uses_custom_var(&self) -> bool {
         let Expr::Name(ast::ExprName {
             id: first_arg_type, ..
@@ -230,3 +236,167 @@ fn is_likely_private_typevar(type_var_name: &str, type_params: Option<&TypeParam
         })
     })
 }
+
+fn add_diagnostic(
+    checker: &mut Checker,
+    method_name: String,
+    type_params: Option<&TypeParams>,
+    parameters: &Parameters,
+    returns: &Expr,
+) {
+    let in_stub = checker.source_type.is_stub();
+
+    let mut diagnostic = Diagnostic::new(
+        CustomTypeVarReturnType {
+            method_name,
+            in_stub,
+        },
+        returns.range(),
+    );
+
+    // See `replace_custom_typevar_with_self`'s doc comment
+    if in_stub {
+        let fix = replace_custom_typevar_with_self(checker, type_params, parameters, returns);
+        diagnostic.set_fix(fix);
+    }
+
+    checker.diagnostics.push(diagnostic);
+}
+
+/// Add a "Replace with `Self`" fix that does the following:
+///
+/// * Import `Self` if necessary
+/// * Remove the first parameter's annotation
+/// * Replace the return annotation with `Self`
+/// * Replace other uses of the original `TypeVar` elsewhere in the signature with `Self`
+/// * Remove that `TypeVar` 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.
+fn replace_custom_typevar_with_self(
+    checker: &Checker,
+    type_params: Option<&TypeParams>,
+    parameters: &Parameters,
+    returns: &Expr,
+) -> Fix {
+    let typevar_name = returns.as_name_expr().unwrap().id();
+
+    let mut all_edits = vec![
+        replace_return_annotation_with_self(returns),
+        remove_first_parameter_annotation(parameters),
+    ];
+
+    if let Some(edit) = import_self(checker, returns.range()) {
+        all_edits.push(edit);
+    }
+
+    if let Some(edit) = remove_typevar_declaration(type_params, typevar_name) {
+        all_edits.push(edit);
+    }
+
+    if let Ok(mut edits) = replace_typevar_usages_with_self(parameters, typevar_name) {
+        all_edits.append(&mut edits);
+    }
+
+    let first = all_edits.remove(0);
+    let rest = all_edits;
+
+    Fix::unsafe_edits(first, rest)
+}
+
+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) {
+        "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 Ok((edit, ..)) = importer.get_or_import_symbol(&request, position, semantic) else {
+        return None;
+    };
+
+    Some(edit)
+}
+
+fn remove_first_parameter_annotation(parameters: &Parameters) -> Edit {
+    let mut non_variadic_positional = parameters.posonlyargs.iter().chain(&parameters.args);
+    let first = &non_variadic_positional.next().unwrap().parameter;
+
+    let name_end = first.name.range.end();
+    let annotation_end = first.range.end();
+
+    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_typevar_usages_with_self(
+    parameters: &Parameters,
+    typevar_name: &str,
+) -> Result<Vec<Edit>, ()> {
+    let mut edits = vec![];
+
+    for parameter in parameters.iter().skip(1) {
+        let Some(annotation) = parameter.annotation() else {
+            continue;
+        };
+        let Expr::Name(name) = annotation else {
+            continue;
+        };
+
+        if name.id.as_str() == typevar_name {
+            let edit = Edit::range_replacement("Self".to_string(), annotation.range());
+            edits.push(edit);
+        } else {
+            // FIXME: Handle complex annotations.
+            return Err(());
+        }
+    }
+
+    Ok(edits)
+}
+
+fn remove_typevar_declaration(type_params: Option<&TypeParams>, name: &str) -> Option<Edit> {
+    let is_declaration_in_question = |type_param: &&TypeParam| -> bool {
+        if let TypeParam::TypeVar(typevar) = type_param {
+            return typevar.name.as_str() == name;
+        };
+
+        false
+    };
+
+    let parameter_list = type_params?;
+    let parameters = &parameter_list.type_params;
+    let first = parameters.first()?;
+
+    if parameter_list.len() == 1 && is_declaration_in_question(&first) {
+        return Some(Edit::range_deletion(parameter_list.range));
+    }
+
+    let (index, ..) = parameters
+        .iter()
+        .find_position(is_declaration_in_question)?;
+
+    let typevar_range = parameters[index].as_type_var().unwrap().range();
+    let last_index = parameters.len() - 1;
+
+    let range = if (0..last_index).contains(&index) {
+        let next_range = parameters[index + 1].range();
+        TextRange::new(typevar_range.start(), next_range.start())
+    } else {
+        let previous_range = parameters[index - 1].range();
+        TextRange::new(previous_range.end(), typevar_range.start())
+    };
+
+    Some(Edit::range_deletion(range))
+}