diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.py b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.py new file mode 100644 index 0000000000000..72bc86ac036f6 --- /dev/null +++ b/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: ... diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.pyi b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.pyi new file mode 100644 index 0000000000000..37dd8dfb574c3 --- /dev/null +++ b/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: ... diff --git a/crates/ruff/src/checkers/ast/analyze/statement.rs b/crates/ruff/src/checkers/ast/analyze/statement.rs index 55e34d28cdb1d..e697d6a476e2d 100644 --- a/crates/ruff/src/checkers/ast/analyze/statement.rs +++ b/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); } diff --git a/crates/ruff/src/codes.rs b/crates/ruff/src/codes.rs index 4bbc0cbc9f541..7874bb8bbb4c1 100644 --- a/crates/ruff/src/codes.rs +++ b/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), diff --git a/crates/ruff/src/rules/flake8_pyi/mod.rs b/crates/ruff/src/rules/flake8_pyi/mod.rs index ec790aec90794..259b02c21a342 100644 --- a/crates/ruff/src/rules/flake8_pyi/mod.rs +++ b/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(()) + } } diff --git a/crates/ruff/src/rules/flake8_pyi/rules/custom_type_var_return_type.rs b/crates/ruff/src/rules/flake8_pyi/rules/custom_type_var_return_type.rs new file mode 100644 index 0000000000000..4ba24634a64ec --- /dev/null +++ b/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) = ¶meter.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) = ¶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: 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 + } + }) + }) +} diff --git a/crates/ruff/src/rules/flake8_pyi/rules/mod.rs b/crates/ruff/src/rules/flake8_pyi/rules/mod.rs index f5e8fda7e710d..e340427d0499f 100644 --- a/crates/ruff/src/rules/flake8_pyi/rules/mod.rs +++ b/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; diff --git a/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.py.snap b/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.py.snap new file mode 100644 index 0000000000000..d1aa2e9116558 --- /dev/null +++ b/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 +--- + diff --git a/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap b/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap new file mode 100644 index 0000000000000..121ee656304d4 --- /dev/null +++ b/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 + | + + diff --git a/ruff.schema.json b/ruff.schema.json index 802011030af2f..6ca8426ed0f35 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -2383,6 +2383,7 @@ "PYI016", "PYI017", "PYI018", + "PYI019", "PYI02", "PYI020", "PYI021",