Minor tweaks
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 ac3b4bb..d5d70b0 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 @@
}
}
-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 @@
&& 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 @@
// 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 @@
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 @@
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!(