diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019.py b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019.py index 23af128a3848d..8e9338574e0ae 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019.py +++ b/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 diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019.pyi index 23af128a3848d..8e9338574e0ae 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI019.pyi +++ b/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 diff --git a/crates/ruff_linter/src/rules/flake8_pyi/mod.rs b/crates/ruff_linter/src/rules/flake8_pyi/mod.rs index 3c700f684031d..72b4f19af8ec5 100644 --- a/crates/ruff_linter/src/rules/flake8_pyi/mod.rs +++ b/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<()> { diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/custom_type_var_return_type.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/custom_type_var_return_type.rs index 900ea4d3fdf6e..426d81543da6e 100644 --- a/crates/ruff_linter/src/rules/flake8_pyi/rules/custom_type_var_return_type.rs +++ b/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. diff --git a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI019_PYI019.py.snap b/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI019_PYI019.py.snap index c40006ed27e7e..eaa6a4bcc533b 100644 --- a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI019_PYI019.py.snap +++ b/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 + | diff --git a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap b/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap index 979abf618fe0a..b3ddd0801d1f1 100644 --- a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap +++ b/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 + |