Skip to content

Commit 5139f76

Browse files
authored
[ty] Infer type of self for decorated methods and properties (#21123)
## Summary Infer a type of unannotated `self` parameters in decorated methods / properties. closes astral-sh/ty#1448 ## Test Plan Existing tests, some new tests.
1 parent aca8ba7 commit 5139f76

File tree

7 files changed

+97
-33
lines changed

7 files changed

+97
-33
lines changed

crates/ruff_benchmark/benches/ty_walltime.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,7 @@ static STATIC_FRAME: Benchmark = Benchmark::new(
226226
max_dep_date: "2025-08-09",
227227
python_version: PythonVersion::PY311,
228228
},
229-
750,
229+
800,
230230
);
231231

232232
#[track_caller]

crates/ty_ide/src/semantic_tokens.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1413,7 +1413,7 @@ u = List.__name__ # __name__ should be variable<CURSOR>
14131413
"property" @ 168..176: Decorator
14141414
"prop" @ 185..189: Method [definition]
14151415
"self" @ 190..194: SelfParameter
1416-
"self" @ 212..216: Variable
1416+
"self" @ 212..216: TypeParameter
14171417
"CONSTANT" @ 217..225: Variable [readonly]
14181418
"obj" @ 227..230: Variable
14191419
"MyClass" @ 233..240: Class

crates/ty_python_semantic/resources/mdtest/annotations/self.md

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ A.implicit_self(1)
116116
Passing `self` implicitly also verifies the type:
117117

118118
```py
119-
from typing import Never
119+
from typing import Never, Callable
120120

121121
class Strange:
122122
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
139139
The name `self` is not special in any way.
140140

141141
```py
142+
def some_decorator(f: Callable) -> Callable:
143+
return f
144+
142145
class B:
143146
def name_does_not_matter(this) -> Self:
144147
reveal_type(this) # revealed: Self@name_does_not_matter
@@ -153,18 +156,45 @@ class B:
153156
reveal_type(self) # revealed: Self@keyword_only
154157
return self
155158

159+
@some_decorator
160+
def decorated_method(self) -> Self:
161+
reveal_type(self) # revealed: Self@decorated_method
162+
return self
163+
156164
@property
157165
def a_property(self) -> Self:
158-
# TODO: Should reveal Self@a_property
159-
reveal_type(self) # revealed: Unknown
166+
reveal_type(self) # revealed: Self@a_property
160167
return self
161168

169+
async def async_method(self) -> Self:
170+
reveal_type(self) # revealed: Self@async_method
171+
return self
172+
173+
@staticmethod
174+
def static_method(self):
175+
# The parameter can be called `self`, but it is not treated as `Self`
176+
reveal_type(self) # revealed: Unknown
177+
178+
@staticmethod
179+
@some_decorator
180+
def decorated_static_method(self):
181+
reveal_type(self) # revealed: Unknown
182+
# TODO: On Python <3.10, this should ideally be rejected, because `staticmethod` objects were not callable.
183+
@some_decorator
184+
@staticmethod
185+
def decorated_static_method_2(self):
186+
reveal_type(self) # revealed: Unknown
187+
162188
reveal_type(B().name_does_not_matter()) # revealed: B
163189
reveal_type(B().positional_only(1)) # revealed: B
164190
reveal_type(B().keyword_only(x=1)) # revealed: B
191+
reveal_type(B().decorated_method()) # revealed: Unknown
165192

166193
# TODO: this should be B
167194
reveal_type(B().a_property) # revealed: Unknown
195+
196+
async def _():
197+
reveal_type(await B().async_method()) # revealed: B
168198
```
169199

170200
This also works for generic classes:

crates/ty_python_semantic/resources/mdtest/overloads.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -598,6 +598,7 @@ class CheckClassMethod:
598598
# error: [invalid-overload]
599599
def try_from3(cls, x: int | str) -> CheckClassMethod | None:
600600
if isinstance(x, int):
601+
# error: [call-non-callable]
601602
return cls(x)
602603
return None
603604

crates/ty_python_semantic/resources/mdtest/snapshots/overloads.md_-_Overloads_-_Invalid_-_Inconsistent_decorat…_-_`@classmethod`_(aaa04d4cfa3adaba).snap

Lines changed: 30 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -53,20 +53,21 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/overloads.md
5353
39 | # error: [invalid-overload]
5454
40 | def try_from3(cls, x: int | str) -> CheckClassMethod | None:
5555
41 | if isinstance(x, int):
56-
42 | return cls(x)
57-
43 | return None
58-
44 |
59-
45 | @overload
60-
46 | @classmethod
61-
47 | def try_from4(cls, x: int) -> CheckClassMethod: ...
62-
48 | @overload
63-
49 | @classmethod
64-
50 | def try_from4(cls, x: str) -> None: ...
65-
51 | @classmethod
66-
52 | def try_from4(cls, x: int | str) -> CheckClassMethod | None:
67-
53 | if isinstance(x, int):
68-
54 | return cls(x)
69-
55 | return None
56+
42 | # error: [call-non-callable]
57+
43 | return cls(x)
58+
44 | return None
59+
45 |
60+
46 | @overload
61+
47 | @classmethod
62+
48 | def try_from4(cls, x: int) -> CheckClassMethod: ...
63+
49 | @overload
64+
50 | @classmethod
65+
51 | def try_from4(cls, x: str) -> None: ...
66+
52 | @classmethod
67+
53 | def try_from4(cls, x: int | str) -> CheckClassMethod | None:
68+
54 | if isinstance(x, int):
69+
55 | return cls(x)
70+
56 | return None
7071
```
7172

7273
# Diagnostics
@@ -124,8 +125,22 @@ error[invalid-overload]: Overloaded function `try_from3` does not use the `@clas
124125
| |
125126
| Missing here
126127
41 | if isinstance(x, int):
127-
42 | return cls(x)
128+
42 | # error: [call-non-callable]
128129
|
129130
info: rule `invalid-overload` is enabled by default
130131
131132
```
133+
134+
```
135+
error[call-non-callable]: Object of type `CheckClassMethod` is not callable
136+
--> src/mdtest_snippet.py:43:20
137+
|
138+
41 | if isinstance(x, int):
139+
42 | # error: [call-non-callable]
140+
43 | return cls(x)
141+
| ^^^^^^
142+
44 | return None
143+
|
144+
info: rule `call-non-callable` is enabled by default
145+
146+
```

crates/ty_python_semantic/src/types/function.rs

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,16 @@ pub struct DataclassTransformerParams<'db> {
185185

186186
impl get_size2::GetSize for DataclassTransformerParams<'_> {}
187187

188+
/// Whether a function should implicitly be treated as a staticmethod based on its name.
189+
pub(crate) fn is_implicit_staticmethod(function_name: &str) -> bool {
190+
matches!(function_name, "__new__")
191+
}
192+
193+
/// Whether a function should implicitly be treated as a classmethod based on its name.
194+
pub(crate) fn is_implicit_classmethod(function_name: &str) -> bool {
195+
matches!(function_name, "__init_subclass__" | "__class_getitem__")
196+
}
197+
188198
/// Representation of a function definition in the AST: either a non-generic function, or a generic
189199
/// function that has not been specialized.
190200
///
@@ -257,17 +267,15 @@ impl<'db> OverloadLiteral<'db> {
257267
/// Returns true if this overload is decorated with `@staticmethod`, or if it is implicitly a
258268
/// staticmethod.
259269
pub(crate) fn is_staticmethod(self, db: &dyn Db) -> bool {
260-
self.has_known_decorator(db, FunctionDecorators::STATICMETHOD) || self.name(db) == "__new__"
270+
self.has_known_decorator(db, FunctionDecorators::STATICMETHOD)
271+
|| is_implicit_staticmethod(self.name(db))
261272
}
262273

263274
/// Returns true if this overload is decorated with `@classmethod`, or if it is implicitly a
264275
/// classmethod.
265276
pub(crate) fn is_classmethod(self, db: &dyn Db) -> bool {
266277
self.has_known_decorator(db, FunctionDecorators::CLASSMETHOD)
267-
|| matches!(
268-
self.name(db).as_str(),
269-
"__init_subclass__" | "__class_getitem__"
270-
)
278+
|| is_implicit_classmethod(self.name(db))
271279
}
272280

273281
fn node<'ast>(

crates/ty_python_semantic/src/types/infer/builder.rs

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ use crate::types::diagnostic::{
7878
};
7979
use crate::types::function::{
8080
FunctionDecorators, FunctionLiteral, FunctionType, KnownFunction, OverloadLiteral,
81+
is_implicit_classmethod, is_implicit_staticmethod,
8182
};
8283
use crate::types::generics::{
8384
GenericContext, InferableTypeVars, LegacyGenericBase, SpecializationBuilder, bind_typevar,
@@ -2580,18 +2581,27 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
25802581
return None;
25812582
}
25822583

2583-
let method = infer_definition_types(db, method_definition)
2584-
.declaration_type(method_definition)
2585-
.inner_type()
2586-
.as_function_literal()?;
2584+
let function_node = function_definition.node(self.module());
2585+
let function_name = &function_node.name;
25872586

2588-
if method.is_classmethod(db) {
2589-
// TODO: set the type for `cls` argument
2590-
return None;
2591-
} else if method.is_staticmethod(db) {
2587+
// TODO: handle implicit type of `cls` for classmethods
2588+
if is_implicit_classmethod(function_name) || is_implicit_staticmethod(function_name) {
25922589
return None;
25932590
}
25942591

2592+
let inference = infer_definition_types(db, method_definition);
2593+
for decorator in &function_node.decorator_list {
2594+
let decorator_ty = inference.expression_type(&decorator.expression);
2595+
if decorator_ty.as_class_literal().is_some_and(|class| {
2596+
matches!(
2597+
class.known(db),
2598+
Some(KnownClass::Classmethod | KnownClass::Staticmethod)
2599+
)
2600+
}) {
2601+
return None;
2602+
}
2603+
}
2604+
25952605
let class_definition = self.index.expect_single_definition(class);
25962606
let class_literal = infer_definition_types(db, class_definition)
25972607
.declaration_type(class_definition)

0 commit comments

Comments
 (0)