[ty] Infer type of self for decorated methods and properties (#21123)

## Summary

Infer a type of unannotated `self` parameters in decorated methods /
properties.

closes https://github.com/astral-sh/ty/issues/1448

## Test Plan

Existing tests, some new tests.
This commit is contained in:
David Peter 2025-10-29 22:22:38 +01:00 committed by GitHub
parent aca8ba76a4
commit 5139f76d1f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 97 additions and 33 deletions

View file

@ -226,7 +226,7 @@ static STATIC_FRAME: Benchmark = Benchmark::new(
max_dep_date: "2025-08-09", max_dep_date: "2025-08-09",
python_version: PythonVersion::PY311, python_version: PythonVersion::PY311,
}, },
750, 800,
); );
#[track_caller] #[track_caller]

View file

@ -1413,7 +1413,7 @@ u = List.__name__ # __name__ should be variable<CURSOR>
"property" @ 168..176: Decorator "property" @ 168..176: Decorator
"prop" @ 185..189: Method [definition] "prop" @ 185..189: Method [definition]
"self" @ 190..194: SelfParameter "self" @ 190..194: SelfParameter
"self" @ 212..216: Variable "self" @ 212..216: TypeParameter
"CONSTANT" @ 217..225: Variable [readonly] "CONSTANT" @ 217..225: Variable [readonly]
"obj" @ 227..230: Variable "obj" @ 227..230: Variable
"MyClass" @ 233..240: Class "MyClass" @ 233..240: Class

View file

@ -116,7 +116,7 @@ A.implicit_self(1)
Passing `self` implicitly also verifies the type: Passing `self` implicitly also verifies the type:
```py ```py
from typing import Never from typing import Never, Callable
class Strange: class Strange:
def can_not_be_called(self: Never) -> None: ... def can_not_be_called(self: Never) -> None: ...
@ -139,6 +139,9 @@ The first parameter of instance methods always has type `Self`, if it is not exp
The name `self` is not special in any way. The name `self` is not special in any way.
```py ```py
def some_decorator(f: Callable) -> Callable:
return f
class B: class B:
def name_does_not_matter(this) -> Self: def name_does_not_matter(this) -> Self:
reveal_type(this) # revealed: Self@name_does_not_matter reveal_type(this) # revealed: Self@name_does_not_matter
@ -153,18 +156,45 @@ class B:
reveal_type(self) # revealed: Self@keyword_only reveal_type(self) # revealed: Self@keyword_only
return self return self
@some_decorator
def decorated_method(self) -> Self:
reveal_type(self) # revealed: Self@decorated_method
return self
@property @property
def a_property(self) -> Self: def a_property(self) -> Self:
# TODO: Should reveal Self@a_property reveal_type(self) # revealed: Self@a_property
reveal_type(self) # revealed: Unknown
return self return self
async def async_method(self) -> Self:
reveal_type(self) # revealed: Self@async_method
return self
@staticmethod
def static_method(self):
# The parameter can be called `self`, but it is not treated as `Self`
reveal_type(self) # revealed: Unknown
@staticmethod
@some_decorator
def decorated_static_method(self):
reveal_type(self) # revealed: Unknown
# TODO: On Python <3.10, this should ideally be rejected, because `staticmethod` objects were not callable.
@some_decorator
@staticmethod
def decorated_static_method_2(self):
reveal_type(self) # revealed: Unknown
reveal_type(B().name_does_not_matter()) # revealed: B reveal_type(B().name_does_not_matter()) # revealed: B
reveal_type(B().positional_only(1)) # revealed: B reveal_type(B().positional_only(1)) # revealed: B
reveal_type(B().keyword_only(x=1)) # revealed: B reveal_type(B().keyword_only(x=1)) # revealed: B
reveal_type(B().decorated_method()) # revealed: Unknown
# TODO: this should be B # TODO: this should be B
reveal_type(B().a_property) # revealed: Unknown reveal_type(B().a_property) # revealed: Unknown
async def _():
reveal_type(await B().async_method()) # revealed: B
``` ```
This also works for generic classes: This also works for generic classes:

View file

@ -598,6 +598,7 @@ class CheckClassMethod:
# error: [invalid-overload] # error: [invalid-overload]
def try_from3(cls, x: int | str) -> CheckClassMethod | None: def try_from3(cls, x: int | str) -> CheckClassMethod | None:
if isinstance(x, int): if isinstance(x, int):
# error: [call-non-callable]
return cls(x) return cls(x)
return None return None

View file

@ -53,20 +53,21 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/overloads.md
39 | # error: [invalid-overload] 39 | # error: [invalid-overload]
40 | def try_from3(cls, x: int | str) -> CheckClassMethod | None: 40 | def try_from3(cls, x: int | str) -> CheckClassMethod | None:
41 | if isinstance(x, int): 41 | if isinstance(x, int):
42 | return cls(x) 42 | # error: [call-non-callable]
43 | return None 43 | return cls(x)
44 | 44 | return None
45 | @overload 45 |
46 | @classmethod 46 | @overload
47 | def try_from4(cls, x: int) -> CheckClassMethod: ... 47 | @classmethod
48 | @overload 48 | def try_from4(cls, x: int) -> CheckClassMethod: ...
49 | @classmethod 49 | @overload
50 | def try_from4(cls, x: str) -> None: ... 50 | @classmethod
51 | @classmethod 51 | def try_from4(cls, x: str) -> None: ...
52 | def try_from4(cls, x: int | str) -> CheckClassMethod | None: 52 | @classmethod
53 | if isinstance(x, int): 53 | def try_from4(cls, x: int | str) -> CheckClassMethod | None:
54 | return cls(x) 54 | if isinstance(x, int):
55 | return None 55 | return cls(x)
56 | return None
``` ```
# Diagnostics # Diagnostics
@ -124,8 +125,22 @@ error[invalid-overload]: Overloaded function `try_from3` does not use the `@clas
| | | |
| Missing here | Missing here
41 | if isinstance(x, int): 41 | if isinstance(x, int):
42 | return cls(x) 42 | # error: [call-non-callable]
| |
info: rule `invalid-overload` is enabled by default info: rule `invalid-overload` is enabled by default
``` ```
```
error[call-non-callable]: Object of type `CheckClassMethod` is not callable
--> src/mdtest_snippet.py:43:20
|
41 | if isinstance(x, int):
42 | # error: [call-non-callable]
43 | return cls(x)
| ^^^^^^
44 | return None
|
info: rule `call-non-callable` is enabled by default
```

View file

@ -185,6 +185,16 @@ pub struct DataclassTransformerParams<'db> {
impl get_size2::GetSize for DataclassTransformerParams<'_> {} impl get_size2::GetSize for DataclassTransformerParams<'_> {}
/// Whether a function should implicitly be treated as a staticmethod based on its name.
pub(crate) fn is_implicit_staticmethod(function_name: &str) -> bool {
matches!(function_name, "__new__")
}
/// Whether a function should implicitly be treated as a classmethod based on its name.
pub(crate) fn is_implicit_classmethod(function_name: &str) -> bool {
matches!(function_name, "__init_subclass__" | "__class_getitem__")
}
/// Representation of a function definition in the AST: either a non-generic function, or a generic /// Representation of a function definition in the AST: either a non-generic function, or a generic
/// function that has not been specialized. /// function that has not been specialized.
/// ///
@ -257,17 +267,15 @@ impl<'db> OverloadLiteral<'db> {
/// Returns true if this overload is decorated with `@staticmethod`, or if it is implicitly a /// Returns true if this overload is decorated with `@staticmethod`, or if it is implicitly a
/// staticmethod. /// staticmethod.
pub(crate) fn is_staticmethod(self, db: &dyn Db) -> bool { pub(crate) fn is_staticmethod(self, db: &dyn Db) -> bool {
self.has_known_decorator(db, FunctionDecorators::STATICMETHOD) || self.name(db) == "__new__" self.has_known_decorator(db, FunctionDecorators::STATICMETHOD)
|| is_implicit_staticmethod(self.name(db))
} }
/// Returns true if this overload is decorated with `@classmethod`, or if it is implicitly a /// Returns true if this overload is decorated with `@classmethod`, or if it is implicitly a
/// classmethod. /// classmethod.
pub(crate) fn is_classmethod(self, db: &dyn Db) -> bool { pub(crate) fn is_classmethod(self, db: &dyn Db) -> bool {
self.has_known_decorator(db, FunctionDecorators::CLASSMETHOD) self.has_known_decorator(db, FunctionDecorators::CLASSMETHOD)
|| matches!( || is_implicit_classmethod(self.name(db))
self.name(db).as_str(),
"__init_subclass__" | "__class_getitem__"
)
} }
fn node<'ast>( fn node<'ast>(

View file

@ -78,6 +78,7 @@ use crate::types::diagnostic::{
}; };
use crate::types::function::{ use crate::types::function::{
FunctionDecorators, FunctionLiteral, FunctionType, KnownFunction, OverloadLiteral, FunctionDecorators, FunctionLiteral, FunctionType, KnownFunction, OverloadLiteral,
is_implicit_classmethod, is_implicit_staticmethod,
}; };
use crate::types::generics::{ use crate::types::generics::{
GenericContext, InferableTypeVars, LegacyGenericBase, SpecializationBuilder, bind_typevar, GenericContext, InferableTypeVars, LegacyGenericBase, SpecializationBuilder, bind_typevar,
@ -2580,17 +2581,26 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
return None; return None;
} }
let method = infer_definition_types(db, method_definition) let function_node = function_definition.node(self.module());
.declaration_type(method_definition) let function_name = &function_node.name;
.inner_type()
.as_function_literal()?;
if method.is_classmethod(db) { // TODO: handle implicit type of `cls` for classmethods
// TODO: set the type for `cls` argument if is_implicit_classmethod(function_name) || is_implicit_staticmethod(function_name) {
return None; return None;
} else if method.is_staticmethod(db) { }
let inference = infer_definition_types(db, method_definition);
for decorator in &function_node.decorator_list {
let decorator_ty = inference.expression_type(&decorator.expression);
if decorator_ty.as_class_literal().is_some_and(|class| {
matches!(
class.known(db),
Some(KnownClass::Classmethod | KnownClass::Staticmethod)
)
}) {
return None; return None;
} }
}
let class_definition = self.index.expect_single_definition(class); let class_definition = self.index.expect_single_definition(class);
let class_literal = infer_definition_types(db, class_definition) let class_literal = infer_definition_types(db, class_definition)