Tweaks

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) {

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,120 +85,128 @@ 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())
     {
         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) = &parameter.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) = &parameter.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
+            }
+        })
+    })
 }