[flake8-pyi] Fix several correctness issues with custom-type-var-return-type (PYI019)

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

@@ -87,3 +87,28 @@ def instance_method_unbound_with_another_parameter[S](self: S, other: S) -> S: .
     def multiple_type_vars[S, *Ts, T](self: S, other: S, /, *args: *Ts, a: T, b: list[T]) -> S: ...
 
     def mixing_old_and_new_style_type_vars[T](self: _S695, a: T, b: T) -> _S695: ...
+
+
+class InvalidButWeDoNotPanic:
+    @classmethod
+    def m[S](cls: type[S], /) -> S[int]: ...
+    def n(self: S) -> S[int]: ...
+
+
+import builtins
+
+class UsesFullyQualifiedType:
+    @classmethod
+    def m[S](cls: builtins.type[S]) -> S: ...  # PYI019
+
+
+def shadowed_type():
+    type = 1
+    class A:
+        @classmethod
+        def m[S](cls: type[S]) -> S: ...  # no error here
+
+
+class SubscriptReturnType:
+    @classmethod
+    def m[S](cls: type[S]) -> type[S]: ...  # PYI019, but no autofix (yet)

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

@@ -87,3 +87,28 @@ class PEP695Fix:
     def multiple_type_vars[S, *Ts, T](self: S, other: S, /, *args: *Ts, a: T, b: list[T]) -> S: ...
 
     def mixing_old_and_new_style_type_vars[T](self: _S695, a: T, b: T) -> _S695: ...
+
+
+class InvalidButWeDoNotPanic:
+    @classmethod
+    def m[S](cls: type[S], /) -> S[int]: ...
+    def n(self: S) -> S[int]: ...
+
+
+import builtins
+
+class UsesFullyQualifiedType:
+    @classmethod
+    def m[S](cls: builtins.type[S]) -> S: ...  # PYI019
+
+
+def shadowed_type():
+    type = 1
+    class A:
+        @classmethod
+        def m[S](cls: type[S]) -> S: ...  # no error here
+
+
+class SubscriptReturnType:
+    @classmethod
+    def m[S](cls: type[S]) -> type[S]: ...  # PYI019, but no autofix (yet)

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

@@ -3,11 +3,12 @@ use crate::importer::ImportRequest;
 use itertools::Itertools;
 use ruff_diagnostics::{Applicability, Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast as ast;
-use ruff_python_ast::helpers::map_subscript;
-use ruff_python_ast::{Expr, Parameters, TypeParam, TypeParams};
+use ruff_python_ast::{
+    self as ast, Expr, ExprName, ExprSubscript, Parameters, TypeParam, TypeParams,
+};
 use ruff_python_semantic::analyze::function_type::{self, FunctionType};
 use ruff_python_semantic::analyze::visibility::{is_abstract, is_overload};
+use ruff_python_semantic::SemanticModel;
 use ruff_text_size::{Ranged, TextRange};
 
 /// ## What it does
@@ -135,7 +136,7 @@ pub(crate) fn custom_type_var_return_type(
         }),
     };
 
-    if method.uses_custom_var() {
+    if method.uses_custom_var(semantic) {
         add_diagnostic(checker, function_def, returns);
     }
 }
@@ -147,9 +148,9 @@ enum Method<'a> {
 }
 
 impl Method<'_> {
-    fn uses_custom_var(&self) -> bool {
+    fn uses_custom_var(&self, semantic: &SemanticModel) -> bool {
         match self {
-            Self::Class(class_method) => class_method.uses_custom_var(),
+            Self::Class(class_method) => class_method.uses_custom_var(semantic),
             Self::Instance(instance_method) => instance_method.uses_custom_var(),
         }
     }
@@ -165,34 +166,45 @@ struct ClassMethod<'a> {
 impl ClassMethod<'_> {
     /// 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 {
+    fn uses_custom_var(&self, semantic: &SemanticModel) -> bool {
+        let Expr::Subscript(ast::ExprSubscript {
+            value: cls_annotation_value,
+            slice: cls_annotation_typevar,
+            ..
+        }) = self.cls_annotation
+        else {
             return false;
         };
 
-        let Expr::Name(value) = value.as_ref() else {
+        let Expr::Name(cls_annotation_typevar) = &**cls_annotation_typevar else {
             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 cls_annotation_typevar = &cls_annotation_typevar.id;
 
-        let Expr::Name(slice) = slice.as_ref() else {
+        if !semantic.match_builtin_expr(cls_annotation_value, "type") {
             return false;
-        };
+        }
 
-        let Expr::Name(return_annotation) = map_subscript(self.returns) else {
-            return false;
+        let return_annotation_typevar = match self.returns {
+            Expr::Name(ExprName { id, .. }) => id,
+            Expr::Subscript(ExprSubscript { value, slice, .. }) => {
+                let Expr::Name(return_annotation_typevar) = &**slice else {
+                    return false;
+                };
+                if !semantic.match_builtin_expr(value, "type") {
+                    return false;
+                }
+                &return_annotation_typevar.id
+            }
+            _ => return false,
         };
 
-        if slice.id != return_annotation.id {
+        if cls_annotation_typevar != return_annotation_typevar {
             return false;
         }
 
-        is_likely_private_typevar(&slice.id, self.type_params)
+        is_likely_private_typevar(cls_annotation_typevar, self.type_params)
     }
 }
 
@@ -216,7 +228,7 @@ impl InstanceMethod<'_> {
 
         let Expr::Name(ast::ExprName {
             id: return_type, ..
-        }) = map_subscript(self.returns)
+        }) = self.returns
         else {
             return false;
         };
@@ -292,9 +304,8 @@ fn replace_custom_typevar_with_self(
         return None;
     }
 
-    // The return annotation is guaranteed to be a name,
-    // as verified by `uses_custom_var()`.
-    let typevar_name = returns.as_name_expr().unwrap().id();
+    // Non-`Name` return annotations are not currently autofixed
+    let typevar_name = &returns.as_name_expr()?.id;
 
     let mut all_edits = vec![
         replace_return_annotation_with_self(returns),

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

@@ -185,3 +185,19 @@ PYI019.py:89:75: PYI019 Methods like `mixing_old_and_new_style_type_vars` should
 89 |     def mixing_old_and_new_style_type_vars[T](self: _S695, a: T, b: T) -> _S695: ...
    |                                                                           ^^^^^ PYI019
    |
+
+PYI019.py:102:40: PYI019 Methods like `m` should return `Self` instead of a custom `TypeVar`
+    |
+100 | class UsesFullyQualifiedType:
+101 |     @classmethod
+102 |     def m[S](cls: builtins.type[S]) -> S: ...  # PYI019
+    |                                        ^ PYI019
+    |
+
+PYI019.py:114:31: PYI019 Methods like `m` should return `Self` instead of a custom `TypeVar`
+    |
+112 | class SubscriptReturnType:
+113 |     @classmethod
+114 |     def m[S](cls: type[S]) -> type[S]: ...  # PYI019, but no autofix (yet)
+    |                               ^^^^^^^ PYI019
+    |

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

@@ -206,3 +206,21 @@ PYI019.pyi:89:75: PYI019 Methods like `mixing_old_and_new_style_type_vars` shoul
    |                                                                           ^^^^^ PYI019
    |
    = help: Replace with `Self`
+
+PYI019.pyi:102:40: PYI019 Methods like `m` should return `Self` instead of a custom `TypeVar`
+    |
+100 | class UsesFullyQualifiedType:
+101 |     @classmethod
+102 |     def m[S](cls: builtins.type[S]) -> S: ...  # PYI019
+    |                                        ^ PYI019
+    |
+    = help: Replace with `Self`
+
+PYI019.pyi:114:31: PYI019 Methods like `m` should return `Self` instead of a custom `TypeVar`
+    |
+112 | class SubscriptReturnType:
+113 |     @classmethod
+114 |     def m[S](cls: type[S]) -> type[S]: ...  # PYI019, but no autofix (yet)
+    |                               ^^^^^^^ PYI019
+    |
+    = help: Replace with `Self`

crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__custom_classmethod_rules_preview.snap

@@ -396,6 +396,7 @@ PYI019.pyi:87:94: PYI019 Methods like `multiple_type_vars` should return `Self`
    87 |+    def multiple_type_vars[*Ts, T](self, other: Self, /, *args: *Ts, a: T, b: list[T]) -> Self: ...
 88 88 | 
 89 89 |     def mixing_old_and_new_style_type_vars[T](self: _S695, a: T, b: T) -> _S695: ...
+90 90 | 
 
 PYI019.pyi:89:75: PYI019 Methods like `mixing_old_and_new_style_type_vars` should return `Self` instead of a custom `TypeVar`
    |
@@ -412,3 +413,34 @@ PYI019.pyi:89:75: PYI019 Methods like `mixing_old_and_new_style_type_vars` shoul
 88 88 | 
 89    |-    def mixing_old_and_new_style_type_vars[T](self: _S695, a: T, b: T) -> _S695: ...
    89 |+    def mixing_old_and_new_style_type_vars[T](self, a: T, b: T) -> Self: ...
+90 90 | 
+91 91 | 
+92 92 | class InvalidButWeDoNotPanic:
+
+PYI019.pyi:102:40: PYI019 [*] Methods like `m` should return `Self` instead of a custom `TypeVar`
+    |
+100 | class UsesFullyQualifiedType:
+101 |     @classmethod
+102 |     def m[S](cls: builtins.type[S]) -> S: ...  # PYI019
+    |                                        ^ PYI019
+    |
+    = help: Replace with `Self`
+
+ℹ Safe fix
+99  99  | 
+100 100 | class UsesFullyQualifiedType:
+101 101 |     @classmethod
+102     |-    def m[S](cls: builtins.type[S]) -> S: ...  # PYI019
+    102 |+    def m(cls) -> Self: ...  # PYI019
+103 103 | 
+104 104 | 
+105 105 | def shadowed_type():
+
+PYI019.pyi:114:31: PYI019 Methods like `m` should return `Self` instead of a custom `TypeVar`
+    |
+112 | class SubscriptReturnType:
+113 |     @classmethod
+114 |     def m[S](cls: type[S]) -> type[S]: ...  # PYI019, but no autofix (yet)
+    |                               ^^^^^^^ PYI019
+    |
+    = help: Replace with `Self`