From 947e1a4e4dd8a0c7fe5dafe2e67a4e10e1e92059 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Fri, 8 Nov 2024 21:06:01 -0500
Subject: [PATCH] Minor tweaks

---
 .../flake8_pyi/rules/non_self_return_type.rs  | 115 ++++++++----------
 1 file changed, 53 insertions(+), 62 deletions(-)

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
index ac3b4bb8d1..d5d70b0250 100644
--- 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
@@ -101,63 +101,6 @@ impl Violation for NonSelfReturnType {
     }
 }
 
-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,
@@ -197,7 +140,7 @@ pub(crate) fn non_self_return_type(
             && is_name(returns, &class_def.name)
             && !is_final(&class_def.decorator_list, semantic)
         {
-            add_diagnostic!(checker, stmt, returns, class_def, name);
+            add_diagnostic(checker, stmt, returns, class_def, name);
         }
         return;
     }
@@ -205,7 +148,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, checker) {
-            add_diagnostic!(checker, stmt, returns, class_def, name);
+            add_diagnostic(checker, stmt, returns, class_def, name);
         }
         return;
     }
@@ -213,7 +156,7 @@ pub(crate) fn non_self_return_type(
     if is_name(returns, &class_def.name) {
         if matches!(name, "__enter__" | "__new__") && !is_final(&class_def.decorator_list, semantic)
         {
-            add_diagnostic!(checker, stmt, returns, class_def, name);
+            add_diagnostic(checker, stmt, returns, class_def, name);
         }
         return;
     }
@@ -223,20 +166,68 @@ pub(crate) fn non_self_return_type(
             if is_iterable_or_iterator(returns, semantic)
                 && subclasses_iterator(class_def, semantic)
             {
-                add_diagnostic!(checker, stmt, returns, class_def, name);
+                add_diagnostic(checker, stmt, returns, class_def, name);
             }
         }
         "__aiter__" => {
             if is_async_iterable_or_iterator(returns, semantic)
                 && subclasses_async_iterator(class_def, semantic)
             {
-                add_diagnostic!(checker, stmt, returns, class_def, name);
+                add_diagnostic(checker, stmt, returns, class_def, name);
             }
         }
         _ => {}
     }
 }
 
+/// Add a diagnostic for the given method.
+fn add_diagnostic(
+    checker: &mut Checker,
+    stmt: &Stmt,
+    returns: &Expr,
+    class_def: &ast::StmtClassDef,
+    method_name: &str,
+) {
+    /// Return an [`Edit`] that imports `typing.Self` from `typing` or `typing_extensions`.
+    fn import_self(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"
+        } else {
+            "typing_extensions"
+        };
+
+        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)
+    }
+
+    /// Generate a [`Fix`] that replaces the return type with `Self`.
+    fn replace_with_self(checker: &mut Checker, range: TextRange) -> Option<Fix> {
+        let import_self = import_self(checker, range)?;
+        let replace_with_self = Edit::range_replacement("Self".to_string(), range);
+        Some(Fix::unsafe_edits(import_self, [replace_with_self]))
+    }
+
+    let mut diagnostic = Diagnostic::new(
+        NonSelfReturnType {
+            class_name: class_def.name.to_string(),
+            method_name: method_name.to_string(),
+        },
+        stmt.identifier(),
+    );
+    if let Some(fix) = replace_with_self(checker, returns.range()) {
+        diagnostic.set_fix(fix);
+    }
+    checker.diagnostics.push(diagnostic);
+}
+
 /// Returns `true` if the method is an in-place binary operator.
 fn is_inplace_bin_op(name: &str) -> bool {
     matches!(