[flake8-pyi] Implement custom_type_var_return_type (PYI019)

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,
             args,
             body,
+            type_params,
             ..
         })
         | Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
@@ -83,6 +84,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             returns,
             args,
             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_typevar_return_type(
+                        checker,
+                        name,
+                        decorator_list,
+                        returns.as_ref().map(AsRef::as_ref),
+                        args,
+                        type_params,
+                    );
+                }
                 if checker.enabled(Rule::StrOrReprDefinedInStub) {
                     flake8_pyi::rules::str_or_repr_defined_in_stub(checker, stmt);
                 }

crates/ruff/src/codes.rs

@@ -636,6 +636,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

@@ -12,6 +12,7 @@ mod tests {
     use crate::registry::Rule;
     use crate::test::test_path;
     use crate::{assert_messages, settings};
+    use crate::settings::types::PythonVersion;
 
     #[test_case(Rule::AnyEqNeAnnotation, Path::new("PYI032.py"))]
     #[test_case(Rule::AnyEqNeAnnotation, Path::new("PYI032.pyi"))]
@@ -110,4 +111,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_typevar_return.rs

@@ -0,0 +1,203 @@
+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_semantic::analyze::visibility::{
+    is_abstract, is_classmethod, is_overload, is_staticmethod,
+};
+use ruff_python_semantic::ScopeKind;
+
+/// ## 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.
+///
+/// ## 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.
+///
+/// ## 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_typevar_return_type(
+    checker: &mut Checker,
+    name: &str,
+    decorator_list: &[Decorator],
+    returns: Option<&Expr>,
+    args: &Arguments,
+    type_params: &[TypeParam],
+) {
+    let ScopeKind::Class(_) = checker.semantic().scope().kind else {
+        return;
+    };
+
+    if args.args.is_empty() && args.posonlyargs.is_empty() {
+        return;
+    }
+
+    let Some(returns) = returns else {
+        return;
+    };
+
+    let return_annotation = if let Expr::Subscript(ast::ExprSubscript { value, .. }) = returns {
+        // Ex) `Type[T]`
+        value
+    } else {
+        // Ex) `Type`, `typing.Type`
+        returns
+    };
+
+    // 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)
+        } 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)
+        };
+
+    if is_violation {
+        checker.diagnostics.push(Diagnostic::new(
+            CustomTypeVarReturnType {
+                method_name: name.to_string(),
+            },
+            return_annotation.range(),
+        ));
+    }
+}
+
+fn check_class_method_for_bad_typevars(
+    checker: &Checker,
+    args: &Arguments,
+    return_annotation: &Expr,
+    type_params: &[TypeParam],
+) -> bool {
+    let ArgWithDefault { def, .. } = &args.args[0];
+
+    let Some(annotation) = &def.annotation 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
+    };
+
+    // 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" {
+        return false;
+    };
+
+    let Expr::Name(ast::ExprName { id: id_slice, .. }) = slice.as_ref() else {
+        return false
+    };
+
+    let Expr::Name(ast::ExprName { id: return_type, .. }) = return_annotation else {
+        return false
+    };
+
+    return_type == id_slice && is_likely_private_typevar(checker, id_slice, type_params)
+}
+
+fn check_instance_method_for_bad_typevars(
+    checker: &Checker,
+    args: &Arguments,
+    return_annotation: &Expr,
+    type_params: &[TypeParam],
+) -> bool {
+    let ArgWithDefault { def, .. } = &args.args[0];
+
+    let Some(annotation) = &def.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(checker, first_arg_type, type_params)
+}
+
+fn is_likely_private_typevar(
+    checker: &Checker,
+    tvar_name: &str,
+    type_params: &[TypeParam],
+) -> bool {
+    if tvar_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
+}

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_typevar_return::*;
 pub(crate) use docstring_in_stubs::*;
 pub(crate) use duplicate_union_member::*;
 pub(crate) use ellipsis_in_non_empty_class_body::*;
@@ -36,6 +37,7 @@ mod bad_version_info_comparison;
 mod collections_named_tuple;
 mod complex_assignment_in_stub;
 mod complex_if_statement_in_stub;
+mod custom_typevar_return;
 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
+---
+