Skip to content

Commit ab46c8d

Browse files
authored
[ty] Add support for properties that return Self (#21335)
## Summary Detect usages of implicit `self` in property getters, which allows us to treat their signature as being generic. closes astral-sh/ty#1502 ## Typing conformance Two new type assertions that are succeeding. ## Ecosystem results Mostly look good. There are a few new false positives related to a bug with constrained typevars that is unrelated to the work here. I reported this as astral-sh/ty#1503. ## Test Plan Added regression tests.
1 parent a6f2dee commit ab46c8d

File tree

3 files changed

+88
-22
lines changed

3 files changed

+88
-22
lines changed

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ 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:
142+
def some_decorator[**P, R](f: Callable[P, R]) -> Callable[P, R]:
143143
return f
144144

145145
class B:
@@ -188,10 +188,10 @@ class B:
188188
reveal_type(B().name_does_not_matter()) # revealed: B
189189
reveal_type(B().positional_only(1)) # revealed: B
190190
reveal_type(B().keyword_only(x=1)) # revealed: B
191+
# TODO: This should deally be `B`
191192
reveal_type(B().decorated_method()) # revealed: Unknown
192193

193-
# TODO: this should be B
194-
reveal_type(B().a_property) # revealed: Unknown
194+
reveal_type(B().a_property) # revealed: B
195195

196196
async def _():
197197
reveal_type(await B().async_method()) # revealed: B

crates/ty_python_semantic/resources/mdtest/properties.md

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,40 @@ c.my_property = 2
4949
c.my_property = "a"
5050
```
5151

52+
## Properties returning `Self`
53+
54+
A property that returns `Self` refers to an instance of the class:
55+
56+
```py
57+
from typing_extensions import Self
58+
59+
class Path:
60+
@property
61+
def parent(self) -> Self:
62+
raise NotImplementedError
63+
64+
reveal_type(Path().parent) # revealed: Path
65+
```
66+
67+
This also works when a setter is defined:
68+
69+
```py
70+
class Node:
71+
@property
72+
def parent(self) -> Self:
73+
raise NotImplementedError
74+
75+
@parent.setter
76+
def parent(self, value: Self) -> None:
77+
pass
78+
79+
root = Node()
80+
child = Node()
81+
child.parent = root
82+
83+
reveal_type(child.parent) # revealed: Node
84+
```
85+
5286
## `property.getter`
5387

5488
`property.getter` can be used to overwrite the getter method of a property. This does not overwrite

crates/ty_python_semantic/src/types/signatures.rs

Lines changed: 51 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,17 @@
1313
use std::{collections::HashMap, slice::Iter};
1414

1515
use itertools::{EitherOrBoth, Itertools};
16+
use ruff_db::parsed::parsed_module;
1617
use ruff_python_ast::ParameterWithDefault;
1718
use smallvec::{SmallVec, smallvec_inline};
1819

1920
use super::{
2021
DynamicType, Type, TypeVarVariance, definition_expression_type, infer_definition_types,
2122
semantic_index,
2223
};
23-
use crate::semantic_index::definition::Definition;
24+
use crate::semantic_index::definition::{Definition, DefinitionKind};
2425
use crate::types::constraints::{ConstraintSet, IteratorConstraintsExtension};
25-
use crate::types::function::FunctionType;
26+
use crate::types::function::{is_implicit_classmethod, is_implicit_staticmethod};
2627
use crate::types::generics::{
2728
GenericContext, InferableTypeVars, typing_self, walk_generic_context,
2829
};
@@ -36,8 +37,11 @@ use crate::{Db, FxOrderSet};
3637
use ruff_python_ast::{self as ast, name::Name};
3738

3839
#[derive(Clone, Copy, Debug)]
40+
#[expect(clippy::struct_excessive_bools)]
3941
struct MethodInformation<'db> {
40-
method: FunctionType<'db>,
42+
is_staticmethod: bool,
43+
is_classmethod: bool,
44+
method_may_be_generic: bool,
4145
class_literal: ClassLiteral<'db>,
4246
class_is_generic: bool,
4347
}
@@ -46,17 +50,49 @@ fn infer_method_information<'db>(
4650
db: &'db dyn Db,
4751
definition: Definition<'db>,
4852
) -> Option<MethodInformation<'db>> {
53+
let DefinitionKind::Function(function_definition) = definition.kind(db) else {
54+
return None;
55+
};
56+
4957
let class_scope_id = definition.scope(db);
5058
let file = class_scope_id.file(db);
59+
let module = parsed_module(db, file).load(db);
5160
let index = semantic_index(db, file);
5261

5362
let class_scope = index.scope(class_scope_id.file_scope_id(db));
5463
let class_node = class_scope.node().as_class()?;
5564

56-
let method = infer_definition_types(db, definition)
57-
.declaration_type(definition)
58-
.inner_type()
59-
.as_function_literal()?;
65+
let function_node = function_definition.node(&module);
66+
let function_name = &function_node.name;
67+
68+
let mut is_staticmethod = is_implicit_classmethod(function_name);
69+
let mut is_classmethod = is_implicit_staticmethod(function_name);
70+
71+
let inference = infer_definition_types(db, definition);
72+
for decorator in &function_node.decorator_list {
73+
let decorator_ty = inference.expression_type(&decorator.expression);
74+
75+
match decorator_ty
76+
.as_class_literal()
77+
.and_then(|class| class.known(db))
78+
{
79+
Some(KnownClass::Staticmethod) => {
80+
is_staticmethod = true;
81+
}
82+
Some(KnownClass::Classmethod) => {
83+
is_classmethod = true;
84+
}
85+
_ => {}
86+
}
87+
}
88+
89+
let method_may_be_generic = match inference.declaration_type(definition).inner_type() {
90+
Type::FunctionLiteral(f) => f.signature(db).overloads.iter().any(|s| {
91+
s.generic_context
92+
.is_some_and(|context| context.variables(db).any(|v| v.typevar(db).is_self(db)))
93+
}),
94+
_ => true,
95+
};
6096

6197
let class_def = index.expect_single_definition(class_node);
6298
let (class_literal, class_is_generic) = match infer_definition_types(db, class_def)
@@ -71,7 +107,9 @@ fn infer_method_information<'db>(
71107
};
72108

73109
Some(MethodInformation {
74-
method,
110+
is_staticmethod,
111+
is_classmethod,
112+
method_may_be_generic,
75113
class_literal,
76114
class_is_generic,
77115
})
@@ -1270,27 +1308,21 @@ impl<'db> Parameters<'db> {
12701308
};
12711309

12721310
let method_info = infer_method_information(db, definition);
1273-
let is_static_or_classmethod = method_info
1274-
.is_some_and(|f| f.method.is_staticmethod(db) || f.method.is_classmethod(db));
1311+
let is_static_or_classmethod =
1312+
method_info.is_some_and(|f| f.is_staticmethod || f.is_classmethod);
12751313

12761314
let inferred_annotation = |arg: &ParameterWithDefault| {
12771315
if let Some(MethodInformation {
1278-
method,
1316+
method_may_be_generic,
12791317
class_literal,
12801318
class_is_generic,
1319+
..
12811320
}) = method_info
12821321
&& !is_static_or_classmethod
12831322
&& arg.parameter.annotation().is_none()
12841323
&& parameters.index(arg.name().id()) == Some(0)
12851324
{
1286-
let method_has_self_in_generic_context =
1287-
method.signature(db).overloads.iter().any(|s| {
1288-
s.generic_context.is_some_and(|context| {
1289-
context.variables(db).any(|v| v.typevar(db).is_self(db))
1290-
})
1291-
});
1292-
1293-
if method_has_self_in_generic_context
1325+
if method_may_be_generic
12941326
|| class_is_generic
12951327
|| class_literal
12961328
.known(db)

0 commit comments

Comments
 (0)