[flake8-pyi] Respect pep8_naming.classmethod-decorators settings when determining if a method is a classmethod in custom-type-var-return-type (PYI019)

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

@@ -44,3 +44,11 @@ def generic_instance_method[S](self: S) -> S: ...  # PYI019
 
 class PEP695GoodDunderNew[T]:
   def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...
+
+
+class CustomClassMethod:
+   # Should be recognised as a classmethod decorator
+   # due to `foo_classmethod being listed in `pep8_naming.classmethod-decorators`
+   # in the settings for this test:
+   @foo_classmethod
+   def foo[S](cls: type[S]) -> S: ...  # PYI019

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

@@ -44,3 +44,11 @@ class PEP695BadDunderNew[T]:
 
 class PEP695GoodDunderNew[T]:
   def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...
+
+
+class CustomClassMethod:
+   # Should be recognised as a classmethod decorator
+   # due to `foo_classmethod being listed in `pep8_naming.classmethod-decorators`
+   # in the settings for this test:
+   @foo_classmethod
+   def foo[S](cls: type[S]) -> S: ...  # PYI019

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

@@ -9,6 +9,7 @@ mod tests {
     use test_case::test_case;
 
     use crate::registry::Rule;
+    use crate::rules::pep8_naming;
     use crate::settings::types::{PreviewMode, PythonVersion};
     use crate::test::test_path;
     use crate::{assert_messages, settings};
@@ -33,8 +34,6 @@ mod tests {
     #[test_case(Rule::ComplexAssignmentInStub, Path::new("PYI017.pyi"))]
     #[test_case(Rule::ComplexIfStatementInStub, Path::new("PYI002.py"))]
     #[test_case(Rule::ComplexIfStatementInStub, Path::new("PYI002.pyi"))]
-    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019.py"))]
-    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019.pyi"))]
     #[test_case(Rule::DocstringInStub, Path::new("PYI021.py"))]
     #[test_case(Rule::DocstringInStub, Path::new("PYI021.pyi"))]
     #[test_case(Rule::DuplicateLiteralMember, Path::new("PYI062.py"))]
@@ -135,6 +134,24 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019.py"))]
+    #[test_case(Rule::CustomTypeVarReturnType, Path::new("PYI019.pyi"))]
+    fn custom_classmethod_rules(rule_code: Rule, path: &Path) -> Result<()> {
+        let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("flake8_pyi").join(path).as_path(),
+            &settings::LinterSettings {
+                pep8_naming: pep8_naming::settings::Settings {
+                    classmethod_decorators: vec!["foo_classmethod".to_string()],
+                    ..pep8_naming::settings::Settings::default()
+                },
+                ..settings::LinterSettings::for_rule(rule_code)
+            },
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
+
     #[test_case(Rule::TypeAliasWithoutAnnotation, Path::new("PYI026.py"))]
     #[test_case(Rule::TypeAliasWithoutAnnotation, Path::new("PYI026.pyi"))]
     fn py38(rule_code: Rule, path: &Path) -> Result<()> {

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

@@ -3,9 +3,8 @@ 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::visibility::{
-    is_abstract, is_classmethod, is_new, is_overload, is_staticmethod,
-};
+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;
@@ -71,7 +70,7 @@ pub(crate) fn custom_type_var_return_type(
     type_params: Option<&TypeParams>,
 ) {
     // Given, e.g., `def foo(self: _S, arg: bytes) -> _T`, extract `_T`.
-    let Some(return_annotation) = returns else {
+    let Some(returns) = returns else {
         return;
     };
 
@@ -86,98 +85,132 @@ pub(crate) fn custom_type_var_return_type(
         return;
     };
 
-    if !checker.semantic().current_scope().kind.is_class() {
-        return;
-    };
+    let semantic = checker.semantic();
 
     // Skip any abstract, static, and overloaded methods.
-    if is_abstract(decorator_list, checker.semantic())
-        || is_overload(decorator_list, checker.semantic())
-        || is_staticmethod(decorator_list, checker.semantic())
-    {
+    if is_abstract(decorator_list, semantic) || is_overload(decorator_list, semantic) {
         return;
     }
 
-    let uses_custom_var: bool =
-        if is_classmethod(decorator_list, checker.semantic()) || is_new(name) {
-            class_method(self_or_cls_annotation, return_annotation, type_params)
-        } else {
-            // If not static, or a class method or __new__ we know it is an instance method
-            instance_method(self_or_cls_annotation, return_annotation, type_params)
-        };
+    let method = match function_type::classify(
+        name,
+        decorator_list,
+        semantic.current_scope(),
+        semantic,
+        &checker.settings.pep8_naming.classmethod_decorators,
+        &checker.settings.pep8_naming.staticmethod_decorators,
+    ) {
+        FunctionType::Function => return,
+        FunctionType::StaticMethod => return,
+        FunctionType::ClassMethod => Method::Class(ClassMethod {
+            cls_annotation: self_or_cls_annotation,
+            returns,
+            type_params,
+        }),
+        FunctionType::Method => Method::Instance(InstanceMethod {
+            self_annotation: self_or_cls_annotation,
+            returns,
+            type_params,
+        }),
+    };
 
-    if uses_custom_var {
+    if method.uses_custom_var() {
         checker.diagnostics.push(Diagnostic::new(
             CustomTypeVarReturnType {
                 method_name: name.to_string(),
             },
-            return_annotation.range(),
+            returns.range(),
         ));
     }
 }
 
-/// Returns `true` if the class method is annotated with a custom `TypeVar` that is likely
-/// private.
-fn class_method(
-    cls_annotation: &Expr,
-    return_annotation: &Expr,
-    type_params: Option<&TypeParams>,
-) -> bool {
-    let Expr::Subscript(ast::ExprSubscript { slice, value, .. }) = cls_annotation else {
-        return false;
-    };
+#[derive(Debug)]
+enum Method<'a> {
+    Class(ClassMethod<'a>),
+    Instance(InstanceMethod<'a>),
+}
 
-    let Expr::Name(value) = value.as_ref() else {
-        return false;
-    };
+impl<'a> Method<'a> {
+    fn uses_custom_var(&self) -> bool {
+        match self {
+            Self::Class(class_method) => class_method.uses_custom_var(),
+            Self::Instance(instance_method) => instance_method.uses_custom_var(),
+        }
+    }
+}
 
-    // Don't error if the first argument is annotated with typing.Type[T].
-    // These are edge cases, and it's hard to give good error messages for them.
-    if value.id != "type" {
-        return false;
-    };
+#[derive(Debug)]
+struct ClassMethod<'a> {
+    cls_annotation: &'a Expr,
+    returns: &'a Expr,
+    type_params: Option<&'a TypeParams>,
+}
 
-    let Expr::Name(slice) = slice.as_ref() else {
-        return false;
-    };
+impl<'a> ClassMethod<'a> {
+    /// 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;
+        };
 
-    let Expr::Name(return_annotation) = map_subscript(return_annotation) else {
-        return false;
-    };
+        let Expr::Name(value) = value.as_ref() else {
+            return false;
+        };
 
-    if slice.id != return_annotation.id {
-        return false;
+        // Don't error if the first argument is annotated with typing.Type[T].
+        // These are edge cases, and it's hard to give good error messages for them.
+        if value.id != "type" {
+            return false;
+        };
+
+        let Expr::Name(slice) = slice.as_ref() else {
+            return false;
+        };
+
+        let Expr::Name(return_annotation) = map_subscript(self.returns) else {
+            return false;
+        };
+
+        if slice.id != return_annotation.id {
+            return false;
+        }
+
+        is_likely_private_typevar(&slice.id, self.type_params)
     }
+}
 
-    is_likely_private_typevar(&slice.id, type_params)
+#[derive(Debug)]
+struct InstanceMethod<'a> {
+    self_annotation: &'a Expr,
+    returns: &'a Expr,
+    type_params: Option<&'a TypeParams>,
 }
 
-/// Returns `true` if the instance method is annotated with a custom `TypeVar` that is likely
-/// private.
-fn instance_method(
-    self_annotation: &Expr,
-    return_annotation: &Expr,
-    type_params: Option<&TypeParams>,
-) -> bool {
-    let Expr::Name(ast::ExprName {
-        id: first_arg_type, ..
-    }) = self_annotation
-    else {
-        return false;
-    };
+impl<'a> InstanceMethod<'a> {
+    /// 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, ..
+        }) = self.self_annotation
+        else {
+            return false;
+        };
 
-    let Expr::Name(ast::ExprName {
-        id: return_type, ..
-    }) = map_subscript(return_annotation)
-    else {
-        return false;
-    };
+        let Expr::Name(ast::ExprName {
+            id: return_type, ..
+        }) = map_subscript(self.returns)
+        else {
+            return false;
+        };
 
-    if first_arg_type != return_type {
-        return false;
-    }
+        if first_arg_type != return_type {
+            return false;
+        }
 
-    is_likely_private_typevar(first_arg_type, type_params)
+        is_likely_private_typevar(first_arg_type, self.type_params)
+    }
 }
 
 /// Returns `true` if the type variable is likely private.

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

@@ -42,4 +42,10 @@ PYI019.py:42:46: PYI019 Methods like `generic_instance_method` should return `ty
    |                                              ^ PYI019
    |
 
-
+PYI019.py:54:32: PYI019 Methods like `foo` should return `typing.Self` instead of a custom `TypeVar`
+   |
+52 |    # in the settings for this test:
+53 |    @foo_classmethod
+54 |    def foo[S](cls: type[S]) -> S: ...  # PYI019
+   |                                ^ PYI019
+   |

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

@@ -42,4 +42,10 @@ PYI019.pyi:42:46: PYI019 Methods like `generic_instance_method` should return `t
    |                                              ^ PYI019
    |
 
-
+PYI019.pyi:54:32: PYI019 Methods like `foo` should return `typing.Self` instead of a custom `TypeVar`
+   |
+52 |    # in the settings for this test:
+53 |    @foo_classmethod
+54 |    def foo[S](cls: type[S]) -> S: ...  # PYI019
+   |                                ^ PYI019
+   |