[flake8-pyi] Implement custom_type_var_return_type (PYI019) (#6204)

## Summary

Implements `Y019` from
[flake8-pyi](https://github.com/PyCQA/flake8-pyi).

The rule checks if

-  instance methods that return `self` 
-  class methods that return an instance of `cls`
- `__new__` methods

Return a custom `TypeVar` instead of `typing.Self` and raises a
violation if this is the case. The rule also covers
[PEP-695](https://peps.python.org/pep-0695/) syntax as introduced in
upstream in https://github.com/PyCQA/flake8-pyi/pull/402

## Test Plan

Added fixtures with test cases from upstream implementation (plus
additional one for an excluded edge case, mentioned in upstream
implementation)

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

@@ -0,0 +1,42 @@
+from typing import TypeVar, Self, Type
+
+_S = TypeVar("_S", bound=BadClass)
+_S2 = TypeVar("_S2", BadClass, GoodClass)
+
+class BadClass:
+    def __new__(cls: type[_S], *args: str, **kwargs: int) -> _S: ...  # Ok
+
+
+    def bad_instance_method(self: _S, arg: bytes) -> _S: ...  # Ok
+
+
+    @classmethod
+    def bad_class_method(cls: type[_S], arg: int) -> _S: ...  # Ok
+
+
+    @classmethod
+    def excluded_edge_case(cls: Type[_S], arg: int) -> _S: ...  # Ok
+
+
+class GoodClass:
+    def __new__(cls: type[Self], *args: list[int], **kwargs: set[str]) -> Self: ...
+    def good_instance_method_1(self: Self, arg: bytes) -> Self: ...
+    def good_instance_method_2(self, arg1: _S2, arg2: _S2) -> _S2: ...
+    @classmethod
+    def good_cls_method_1(cls: type[Self], arg: int) -> Self: ...
+    @classmethod
+    def good_cls_method_2(cls, arg1: _S, arg2: _S) -> _S: ...
+    @staticmethod
+    def static_method(arg1: _S) -> _S: ...
+
+
+# Python > 3.12
+class PEP695BadDunderNew[T]:
+  def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ...  # Ok
+
+
+  def generic_instance_method[S](self: S) -> S: ...  # Ok
+
+
+class PEP695GoodDunderNew[T]:
+  def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...

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

@@ -0,0 +1,42 @@
+from typing import TypeVar, Self, Type
+
+_S = TypeVar("_S", bound=BadClass)
+_S2 = TypeVar("_S2", BadClass, GoodClass)
+
+class BadClass:
+    def __new__(cls: type[_S], *args: str, **kwargs: int) -> _S: ...  # PYI019
+
+
+    def bad_instance_method(self: _S, arg: bytes) -> _S: ...  # PYI019
+
+
+    @classmethod
+    def bad_class_method(cls: type[_S], arg: int) -> _S: ...  # PYI019
+
+
+    @classmethod
+    def excluded_edge_case(cls: Type[_S], arg: int) -> _S: ...  # Ok
+
+
+class GoodClass:
+    def __new__(cls: type[Self], *args: list[int], **kwargs: set[str]) -> Self: ...
+    def good_instance_method_1(self: Self, arg: bytes) -> Self: ...
+    def good_instance_method_2(self, arg1: _S2, arg2: _S2) -> _S2: ...
+    @classmethod
+    def good_cls_method_1(cls: type[Self], arg: int) -> Self: ...
+    @classmethod
+    def good_cls_method_2(cls, arg1: _S, arg2: _S) -> _S: ...
+    @staticmethod
+    def static_method(arg1: _S) -> _S: ...
+
+
+# Python > 3.12
+class PEP695BadDunderNew[T]:
+  def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ...  # PYI019
+
+
+  def generic_instance_method[S](self: S) -> S: ...  # PYI019
+
+
+class PEP695GoodDunderNew[T]:
+  def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...

crates/ruff/src/checkers/ast/analyze/statement.rs

@@ -75,6 +75,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             returns,
             parameters,
             body,
+            type_params,
             ..
         })
         | Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
@@ -83,6 +84,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             returns,
             parameters,
             body,
+            type_params,
             ..
         }) => {
             if checker.enabled(Rule::DjangoNonLeadingReceiverDecorator) {
@@ -155,6 +157,16 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                         stmt.is_async_function_def_stmt(),
                     );
                 }
+                if checker.enabled(Rule::CustomTypeVarReturnType) {
+                    flake8_pyi::rules::custom_type_var_return_type(
+                        checker,
+                        name,
+                        decorator_list,
+                        returns.as_ref().map(AsRef::as_ref),
+                        parameters,
+                        type_params.as_ref(),
+                    );
+                }
                 if checker.enabled(Rule::StrOrReprDefinedInStub) {
                     flake8_pyi::rules::str_or_repr_defined_in_stub(checker, stmt);
                 }

crates/ruff/src/codes.rs

@@ -637,6 +637,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Flake8Pyi, "016") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::DuplicateUnionMember),
         (Flake8Pyi, "017") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::ComplexAssignmentInStub),
         (Flake8Pyi, "018") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::UnusedPrivateTypeVar),
+        (Flake8Pyi, "019") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::CustomTypeVarReturnType),
         (Flake8Pyi, "020") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::QuotedAnnotationInStub),
         (Flake8Pyi, "021") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::DocstringInStub),
         (Flake8Pyi, "024") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::CollectionsNamedTuple),

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

@@ -10,6 +10,7 @@ mod tests {
     use test_case::test_case;
 
     use crate::registry::Rule;
+    use crate::settings::types::PythonVersion;
     use crate::test::test_path;
     use crate::{assert_messages, settings};
 
@@ -112,4 +113,19 @@ mod tests {
         assert_messages!(snapshot, diagnostics);
         Ok(())
     }
+
+    #[test_case(Path::new("PYI019.py"))]
+    #[test_case(Path::new("PYI019.pyi"))]
+    fn custom_type_var_return_type(path: &Path) -> Result<()> {
+        let snapshot = format!("{}_{}", "PYI019", path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("flake8_pyi").join(path).as_path(),
+            &settings::Settings {
+                target_version: PythonVersion::Py312,
+                ..settings::Settings::for_rules(vec![Rule::CustomTypeVarReturnType])
+            },
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
 }

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

@@ -0,0 +1,212 @@
+use ruff_diagnostics::{Diagnostic, Violation};
+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, ParameterWithDefault, Parameters, Ranged, TypeParam, TypeParams,
+};
+use ruff_python_semantic::analyze::visibility::{
+    is_abstract, is_classmethod, is_new, is_overload, is_staticmethod,
+};
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// 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.
+///
+/// This check currently applies to instance methods that return `self`, class
+/// methods that return an instance of `cls`, and `__new__` methods.
+///
+/// ## Example
+/// ```python
+/// class Foo:
+///     def __new__(cls: type[_S], *args: str, **kwargs: int) -> _S:
+///         ...
+///
+///     def foo(self: _S, arg: bytes) -> _S:
+///         ...
+///
+///     @classmethod
+///     def bar(cls: type[_S], arg: int) -> _S:
+///         ...
+/// ```
+///
+/// Use instead:
+/// ```python
+/// from typing import Self
+///
+///
+/// class Foo:
+///     def __new__(cls: type[Self], *args: str, **kwargs: int) -> Self:
+///         ...
+///
+///     def foo(self: Self, arg: bytes) -> Self:
+///         ...
+///
+///     @classmethod
+///     def bar(cls: type[Self], arg: int) -> Self:
+///         ...
+/// ```
+#[violation]
+pub struct CustomTypeVarReturnType {
+    method_name: String,
+}
+
+impl Violation for CustomTypeVarReturnType {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let CustomTypeVarReturnType { method_name } = self;
+        format!(
+            "Methods like `{method_name}` should return `typing.Self` instead of a custom `TypeVar`"
+        )
+    }
+}
+
+/// PYI019
+pub(crate) fn custom_type_var_return_type(
+    checker: &mut Checker,
+    name: &str,
+    decorator_list: &[Decorator],
+    returns: Option<&Expr>,
+    args: &Parameters,
+    type_params: Option<&TypeParams>,
+) {
+    if args.args.is_empty() && args.posonlyargs.is_empty() {
+        return;
+    }
+
+    let Some(returns) = returns else {
+        return;
+    };
+
+    if !checker.semantic().scope().kind.is_class() {
+        return;
+    };
+
+    // 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 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
+            instance_method(args, returns, type_params)
+        };
+
+    if uses_custom_var {
+        checker.diagnostics.push(Diagnostic::new(
+            CustomTypeVarReturnType {
+                method_name: name.to_string(),
+            },
+            returns.range(),
+        ));
+    }
+}
+
+/// 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: Option<&TypeParams>,
+) -> bool {
+    let ParameterWithDefault { parameter, .. } = &args.args[0];
+
+    let Some(annotation) = &parameter.annotation else {
+        return false;
+    };
+
+    let Expr::Subscript(ast::ExprSubscript { slice, value, .. }) = annotation.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].
+    // 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) = return_annotation else {
+        return false;
+    };
+
+    if slice.id != return_annotation.id {
+        return false;
+    }
+
+    is_likely_private_typevar(&slice.id, type_params)
+}
+
+/// 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: Option<&TypeParams>,
+) -> bool {
+    let ParameterWithDefault { parameter, .. } = &args.args[0];
+
+    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: return_type, ..
+    }) = return_annotation
+    else {
+        return false;
+    };
+
+    if first_arg_type != return_type {
+        return false;
+    }
+
+    is_likely_private_typevar(first_arg_type, type_params)
+}
+
+/// 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;
+    }
+    // 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
+            }
+        })
+    })
+}

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

@@ -3,6 +3,7 @@ pub(crate) use bad_version_info_comparison::*;
 pub(crate) use collections_named_tuple::*;
 pub(crate) use complex_assignment_in_stub::*;
 pub(crate) use complex_if_statement_in_stub::*;
+pub(crate) use custom_type_var_return_type::*;
 pub(crate) use docstring_in_stubs::*;
 pub(crate) use duplicate_union_member::*;
 pub(crate) use ellipsis_in_non_empty_class_body::*;
@@ -37,6 +38,7 @@ mod bad_version_info_comparison;
 mod collections_named_tuple;
 mod complex_assignment_in_stub;
 mod complex_if_statement_in_stub;
+mod custom_type_var_return_type;
 mod docstring_in_stubs;
 mod duplicate_union_member;
 mod ellipsis_in_non_empty_class_body;

crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff/src/rules/flake8_pyi/mod.rs
+---
+

crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap

@@ -0,0 +1,38 @@
+---
+source: crates/ruff/src/rules/flake8_pyi/mod.rs
+---
+PYI019.pyi:7:62: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
+  |
+6 | class BadClass:
+7 |     def __new__(cls: type[_S], *args: str, **kwargs: int) -> _S: ...  # PYI019
+  |                                                              ^^ PYI019
+  |
+
+PYI019.pyi:10:54: PYI019 Methods like `bad_instance_method` should return `typing.Self` instead of a custom `TypeVar`
+   |
+10 |     def bad_instance_method(self: _S, arg: bytes) -> _S: ...  # PYI019
+   |                                                      ^^ PYI019
+   |
+
+PYI019.pyi:14:54: PYI019 Methods like `bad_class_method` should return `typing.Self` instead of a custom `TypeVar`
+   |
+13 |     @classmethod
+14 |     def bad_class_method(cls: type[_S], arg: int) -> _S: ...  # PYI019
+   |                                                      ^^ PYI019
+   |
+
+PYI019.pyi:35:63: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
+   |
+33 | # Python > 3.12
+34 | class PEP695BadDunderNew[T]:
+35 |   def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ...  # PYI019
+   |                                                               ^ PYI019
+   |
+
+PYI019.pyi:38:46: PYI019 Methods like `generic_instance_method` should return `typing.Self` instead of a custom `TypeVar`
+   |
+38 |   def generic_instance_method[S](self: S) -> S: ...  # PYI019
+   |                                              ^ PYI019
+   |
+
+

ruff.schema.json

@@ -2383,6 +2383,7 @@
         "PYI016",
         "PYI017",
         "PYI018",
+        "PYI019",
         "PYI02",
         "PYI020",
         "PYI021",