diff --git a/crates/ruff/src/checkers/ast/analyze/statement.rs b/crates/ruff/src/checkers/ast/analyze/statement.rs index 12fdf21ad5ecd8..146d3b4382df53 100644 --- a/crates/ruff/src/checkers/ast/analyze/statement.rs +++ b/crates/ruff/src/checkers/ast/analyze/statement.rs @@ -163,8 +163,8 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) { name, decorator_list, returns.as_ref().map(AsRef::as_ref), - args, - type_params, + parameters, + type_params.as_ref(), ); } if checker.enabled(Rule::StrOrReprDefinedInStub) { diff --git a/crates/ruff/src/rules/flake8_pyi/rules/custom_typevar_return.rs b/crates/ruff/src/rules/flake8_pyi/rules/custom_typevar_return.rs index 8b5ebda1343c3e..c26367d7a2607e 100644 --- a/crates/ruff/src/rules/flake8_pyi/rules/custom_typevar_return.rs +++ b/crates/ruff/src/rules/flake8_pyi/rules/custom_typevar_return.rs @@ -1,22 +1,27 @@ -use crate::checkers::ast::Checker; -use crate::settings::types::PythonVersion; use ruff_diagnostics::{Diagnostic, Violation}; use ruff_macros::{derive_message_formats, violation}; use ruff_python_ast as ast; -use ruff_python_ast::{ArgWithDefault, Arguments, Decorator, Expr, Ranged, TypeParam}; +use ruff_python_ast::helpers::map_subscript; +use ruff_python_ast::{ + Decorator, Expr, ParameterWithDefault, Parameters, Ranged, TypeParam, TypeParams, +}; use ruff_python_semantic::analyze::visibility::{ - is_abstract, is_classmethod, is_overload, is_staticmethod, + is_abstract, is_classmethod, is_new, is_overload, is_staticmethod, }; -use ruff_python_semantic::ScopeKind; + +use crate::checkers::ast::Checker; /// ## What it does -/// Checks if certain methods define a custom `TypeVar`s for their return annotation instead of -/// using `typing_extensions.Self`. This check currently applies for instance methods that return -/// `self`, class methods that return an instance of `cls`, and `__new__` methods. +/// Checks for methods that define a custom `TypeVar` for their return type +/// annotation instead of using `typing_extensions.Self`. /// /// ## Why is this bad? -/// If certain methods are annotated with a custom `TypeVar` type, and the class is subclassed, -/// type checkers will not be able to infer the correct return type. +/// If certain methods are annotated with a custom `TypeVar` type, and the +/// class is subclassed, type checkers will not be able to infer the correct +/// return type. +/// +/// This check currently applies to instance methods that return `self`, class +/// methods that return an instance of `cls`, and `__new__` methods. /// /// ## Example /// ```python @@ -58,7 +63,7 @@ impl Violation for CustomTypeVarReturnType { fn message(&self) -> String { let CustomTypeVarReturnType { method_name } = self; format!( - "Methods like {method_name} should return `typing.Self` instead of a custom TypeVar" + "Methods like `{method_name}` should return `typing.Self` instead of a custom `TypeVar`" ) } } @@ -69,13 +74,9 @@ pub(crate) fn custom_typevar_return_type( name: &str, decorator_list: &[Decorator], returns: Option<&Expr>, - args: &Arguments, - type_params: &[TypeParam], + args: &Parameters, + type_params: Option<&TypeParams>, ) { - let ScopeKind::Class(_) = checker.semantic().scope().kind else { - return; - }; - if args.args.is_empty() && args.posonlyargs.is_empty() { return; } @@ -84,15 +85,11 @@ pub(crate) fn custom_typevar_return_type( return; }; - let return_annotation = if let Expr::Subscript(ast::ExprSubscript { value, .. }) = returns { - // Ex) `Type[T]` - value - } else { - // Ex) `Type`, `typing.Type` - returns + if !checker.semantic().scope().kind.is_class() { + return; }; - // Skip any abstract, static and overloaded methods. + // 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()) @@ -100,104 +97,116 @@ pub(crate) fn custom_typevar_return_type( return; } - let is_violation: bool = - if is_classmethod(decorator_list, checker.semantic()) || name == "__new__" { - check_class_method_for_bad_typevars(checker, args, return_annotation, type_params) + let returns = map_subscript(returns); + + let uses_custom_var: bool = + if is_classmethod(decorator_list, checker.semantic()) || is_new(name) { + class_method(args, returns, type_params) } else { // If not static, or a class method or __new__ we know it is an instance method - check_instance_method_for_bad_typevars(checker, args, return_annotation, type_params) + instance_method(args, returns, type_params) }; - if is_violation { + if uses_custom_var { checker.diagnostics.push(Diagnostic::new( CustomTypeVarReturnType { method_name: name.to_string(), }, - return_annotation.range(), + returns.range(), )); } } -fn check_class_method_for_bad_typevars( - checker: &Checker, - args: &Arguments, +/// Returns `true` if the class method is annotated with a custom `TypeVar` that is likely +/// private. +fn class_method( + args: &Parameters, return_annotation: &Expr, - type_params: &[TypeParam], + type_params: Option<&TypeParams>, ) -> bool { - let ArgWithDefault { def, .. } = &args.args[0]; + let ParameterWithDefault { parameter, .. } = &args.args[0]; - let Some(annotation) = &def.annotation else { - return false + let Some(annotation) = ¶meter.annotation else { + return false; }; - let Expr::Subscript(ast::ExprSubscript{slice, value, ..}) = annotation.as_ref() else { - return false + let Expr::Subscript(ast::ExprSubscript { slice, value, .. }) = annotation.as_ref() else { + return false; }; - let Expr::Name(ast::ExprName{ id: id_value, .. }) = value.as_ref() else { - return false + let Expr::Name(value) = value.as_ref() else { + return false; }; - // Don't error if the first argument is annotated with typing.Type[T] + // 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 id_value != "type" { + if value.id != "type" { return false; }; - let Expr::Name(ast::ExprName { id: id_slice, .. }) = slice.as_ref() else { - return false + let Expr::Name(slice) = slice.as_ref() else { + return false; }; - let Expr::Name(ast::ExprName { id: return_type, .. }) = return_annotation else { - return false + let Expr::Name(return_annotation) = return_annotation else { + return false; }; - return_type == id_slice && is_likely_private_typevar(checker, id_slice, type_params) + if slice.id != return_annotation.id { + return false; + } + + is_likely_private_typevar(&slice.id, type_params) } -fn check_instance_method_for_bad_typevars( - checker: &Checker, - args: &Arguments, +/// Returns `true` if the instance method is annotated with a custom `TypeVar` that is likely +/// private. +fn instance_method( + args: &Parameters, return_annotation: &Expr, - type_params: &[TypeParam], + type_params: Option<&TypeParams>, ) -> bool { - let ArgWithDefault { def, .. } = &args.args[0]; + let ParameterWithDefault { parameter, .. } = &args.args[0]; - let Some(annotation) = &def.annotation else { - return false + let Some(annotation) = ¶meter.annotation else { + return false; }; - let Expr::Name(ast::ExprName{id: first_arg_type,..}) = annotation.as_ref() else { - return false + let Expr::Name(ast::ExprName { + id: first_arg_type, .. + }) = annotation.as_ref() + else { + return false; }; - let Expr::Name(ast::ExprName { id: return_type, .. }) = return_annotation else { - return false + let Expr::Name(ast::ExprName { + id: return_type, .. + }) = return_annotation + else { + return false; }; if first_arg_type != return_type { return false; } - is_likely_private_typevar(checker, first_arg_type, type_params) + is_likely_private_typevar(first_arg_type, type_params) } -fn is_likely_private_typevar( - checker: &Checker, - tvar_name: &str, - type_params: &[TypeParam], -) -> bool { - if tvar_name.starts_with('_') { +/// Returns `true` if the type variable is likely private. +fn is_likely_private_typevar(type_var_name: &str, type_params: Option<&TypeParams>) -> bool { + // Ex) `_T` + if type_var_name.starts_with('_') { return true; } - if checker.settings.target_version < PythonVersion::Py312 { - return false; - } - - for param in type_params { - if let TypeParam::TypeVar(ast::TypeParamTypeVar { name, .. }) = param { - return name == tvar_name; - } - } - false + // Ex) `class Foo[T]: ...` + type_params.is_some_and(|type_params| { + type_params.iter().any(|type_param| { + if let TypeParam::TypeVar(ast::TypeParamTypeVar { name, .. }) = type_param { + name == type_var_name + } else { + false + } + }) + }) }