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

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

@@ -1,14 +1,15 @@
 use ruff_python_ast::{self as ast, Decorator, Expr, Parameters, Stmt};
 
-use ruff_diagnostics::{Diagnostic, Violation};
+use crate::checkers::ast::Checker;
+use crate::importer::ImportRequest;
+use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::map_subscript;
 use ruff_python_ast::identifier::Identifier;
 use ruff_python_semantic::analyze;
 use ruff_python_semantic::analyze::visibility::{is_abstract, is_final, is_overload};
 use ruff_python_semantic::{ScopeKind, SemanticModel};
-
-use crate::checkers::ast::Checker;
+use ruff_text_size::{Ranged, TextRange};
 
 /// ## What it does
 /// Checks for methods that are annotated with a fixed return type which
@@ -79,12 +80,15 @@ pub struct NonSelfReturnType {
 }
 
 impl Violation for NonSelfReturnType {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
         let NonSelfReturnType {
             class_name,
             method_name,
         } = self;
+
         if matches!(class_name.as_str(), "__new__") {
             "`__new__` methods usually return `self` at runtime".to_string()
         } else {
@@ -93,10 +97,67 @@ impl Violation for NonSelfReturnType {
     }
 
     fn fix_title(&self) -> Option<String> {
-        Some("Consider using `typing_extensions.Self` as return type".to_string())
+        Some("Use `Self` as return type".to_string())
     }
 }
 
+fn import_self_edit(checker: &Checker, range: TextRange) -> Option<Edit> {
+    let target_version = checker.settings.target_version.as_tuple();
+    let source_module = if checker.source_type.is_stub() || target_version < (3, 11) {
+        "typing_extensions"
+    } else {
+        "typing"
+    };
+
+    let (importer, semantic) = (checker.importer(), checker.semantic());
+    let request = ImportRequest::import_from(source_module, "Self");
+
+    let Ok((edit, ..)) = importer.get_or_import_symbol(&request, range.start(), semantic) else {
+        return None;
+    };
+
+    Some(edit)
+}
+
+fn replace_with_self_fix(checker: &mut Checker, range: TextRange) -> Option<Fix> {
+    let import_self = import_self_edit(checker, range)?;
+    let replace_with_self = Edit::range_replacement("Self".to_string(), range);
+
+    Some(Fix::unsafe_edits(import_self, [replace_with_self]))
+}
+
+fn add_diagnostic(
+    checker: &mut Checker,
+    stmt: &Stmt,
+    returns: &Expr,
+    class_name: String,
+    method_name: String,
+) {
+    let kind = NonSelfReturnType {
+        class_name,
+        method_name,
+    };
+    let mut diagnostic = Diagnostic::new(kind, stmt.identifier());
+
+    if let Some(fix) = replace_with_self_fix(checker, returns.range()) {
+        diagnostic.set_fix(fix);
+    }
+
+    checker.diagnostics.push(diagnostic);
+}
+
+macro_rules! add_diagnostic {
+    ($checker:ident, $stmt:ident, $returns:ident, $class_def:ident, $method_name:ident) => {
+        add_diagnostic(
+            $checker,
+            $stmt,
+            $returns,
+            $class_def.name.to_string(),
+            $method_name.to_string(),
+        );
+    };
+}
+
 /// PYI034
 pub(crate) fn non_self_return_type(
     checker: &mut Checker,
@@ -136,41 +197,23 @@ pub(crate) fn non_self_return_type(
             && is_name(returns, &class_def.name)
             && !is_final(&class_def.decorator_list, semantic)
         {
-            checker.diagnostics.push(Diagnostic::new(
-                NonSelfReturnType {
-                    class_name: class_def.name.to_string(),
-                    method_name: name.to_string(),
-                },
-                stmt.identifier(),
-            ));
+            add_diagnostic!(checker, stmt, returns, class_def, name);
         }
         return;
     }
 
     // In-place methods that are expected to return `Self`.
     if is_inplace_bin_op(name) {
         if !is_self(returns, checker) {
-            checker.diagnostics.push(Diagnostic::new(
-                NonSelfReturnType {
-                    class_name: class_def.name.to_string(),
-                    method_name: name.to_string(),
-                },
-                stmt.identifier(),
-            ));
+            add_diagnostic!(checker, stmt, returns, class_def, name);
         }
         return;
     }
 
     if is_name(returns, &class_def.name) {
         if matches!(name, "__enter__" | "__new__") && !is_final(&class_def.decorator_list, semantic)
         {
-            checker.diagnostics.push(Diagnostic::new(
-                NonSelfReturnType {
-                    class_name: class_def.name.to_string(),
-                    method_name: name.to_string(),
-                },
-                stmt.identifier(),
-            ));
+            add_diagnostic!(checker, stmt, returns, class_def, name);
         }
         return;
     }
@@ -180,26 +223,14 @@ pub(crate) fn non_self_return_type(
             if is_iterable_or_iterator(returns, semantic)
                 && subclasses_iterator(class_def, semantic)
             {
-                checker.diagnostics.push(Diagnostic::new(
-                    NonSelfReturnType {
-                        class_name: class_def.name.to_string(),
-                        method_name: name.to_string(),
-                    },
-                    stmt.identifier(),
-                ));
+                add_diagnostic!(checker, stmt, returns, class_def, name);
             }
         }
         "__aiter__" => {
             if is_async_iterable_or_iterator(returns, semantic)
                 && subclasses_async_iterator(class_def, semantic)
             {
-                checker.diagnostics.push(Diagnostic::new(
-                    NonSelfReturnType {
-                        class_name: class_def.name.to_string(),
-                        method_name: name.to_string(),
-                    },
-                    stmt.identifier(),
-                ));
+                add_diagnostic!(checker, stmt, returns, class_def, name);
             }
         }
         _ => {}