Visitar URL original
[ty] Generalize nested types in inferred return types (#29488) · astral-sh/ruff@59a9e85 · GitHub
Skip to content

Commit 59a9e85

Browse files
sharkdpsharkdp
authored
[ty] Generalize nested types in inferred return types (#29488)
## Summary Generalize/promote types in inferred return types that refer to local nested definitions. For example, consider this unannotated decorator: ```python def register(event: str): def wrapper(func): # ... return func return wrapper ``` We do not want to infer `def wrapper(func) -> Unknown` as the return type of `register`. This would not just expose internal details, it would also incorrectly suggest that each call to `register` returns the same function object. Instead, we generalize to an anonymous `(func) -> Unknown` callable. For nested classes, we cannot easily do the same, so we return `Unknown` instead. In the future, to improve this, we might want to synthesize a protocol that has the same attributes and methods as the nested class. ## Test strategy Updated existing mdtests
1 parent 0d562aa commit 59a9e85

6 files changed

Lines changed: 105 additions & 25 deletions

File tree

‎crates/ty_ide/src/inlay_hints.rs‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1082,11 +1082,9 @@ mod tests {
10821082
return 42
10831083
"#,
10841084
);
1085-
// TODO: Do not expose `def wrapper` in the return-type hint for `stringify`,
1086-
// and ideally show the return type of `str` for the returned callable.
10871085
assert_snapshot!(test.inlay_hints(), @r#"
10881086
1089-
def stringify(function)[ -> def wrapper() -> Unknown]:
1087+
def stringify(function)[ -> () -> str]:
10901088
def wrapper()[ -> str]:
10911089
return str([object=]function())
10921090
@@ -1098,15 +1096,15 @@ mod tests {
10981096
10991097
---------------------------------------------
11001098
info[inlay-hint-location]: Inlay Hint Target
1101-
--> stdlib/ty_extensions/_internal.pyi:LL:1
1099+
--> stdlib/builtins.pyi:LL:7
11021100
|
1103-
LL | Unknown: _SpecialForm
1104-
| ^^^^^^^
1101+
LL | class str(Sequence[str]):
1102+
| ^^^
11051103
info: Source
1106-
--> main2.py:LL:46
1104+
--> main2.py:LL:35
11071105
|
1108-
LL | def stringify(function)[ -> def wrapper() -> Unknown]:
1109-
| ^^^^^^^
1106+
LL | def stringify(function)[ -> () -> str]:
1107+
| ^^^
11101108
11111109
info[inlay-hint-location]: Inlay Hint Target
11121110
--> stdlib/builtins.pyi:LL:7

‎crates/ty_python_semantic/resources/mdtest/function/return_type_inference.md‎

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -299,8 +299,7 @@ def returns_class():
299299
Inner1 = returns_class()
300300
Inner2 = returns_class()
301301

302-
# TODO: This should be `bool` or `Literal[False]`, not `Literal[True]`
303-
reveal_type(Inner1 is Inner2) # revealed: Literal[True]
302+
reveal_type(Inner1 is Inner2) # revealed: bool
304303

305304
# TODO: This should ideally be an error
306305
i: Inner1 = Inner2()
@@ -310,15 +309,13 @@ Therefore, the return type of `returns_class` cannot be `<class 'Inner'>`. Inste
310309
just treat it as `Unknown`.
311310

312311
```py
313-
# TODO: Should be `Unknown` or maybe `type`
314-
reveal_type(Inner1) # revealed: <class 'Inner'>
312+
reveal_type(Inner1) # revealed: Unknown
315313
```
316314

317315
Unfortunately, this means that we cannot infer precise types for attributes/methods on those types:
318316

319317
```py
320-
# TODO: Should be `Unknown`
321-
reveal_type(Inner1().attr) # revealed: int
318+
reveal_type(Inner1().attr) # revealed: Unknown
322319
```
323320

324321
If we wanted to improve this, we could return a synthetic protocol representing the structure of the
@@ -342,11 +339,9 @@ def returns_callable():
342339
callable1 = returns_callable()
343340
callable2 = returns_callable()
344341

345-
# TODO: This should be `bool` or `Literal[False]`, not `Literal[True]`
346-
reveal_type(callable1 is callable2) # revealed: Literal[True]
342+
reveal_type(callable1 is callable2) # revealed: bool
347343

348-
# TODO: should be `Callable[[], int]`
349-
reveal_type(callable1) # revealed: def inner() -> int
344+
reveal_type(callable1) # revealed: () -> int
350345
reveal_type(callable1()) # revealed: int
351346
```
352347

@@ -369,12 +364,11 @@ def returns_instance[T](x: T):
369364
def _(x: int):
370365
inner = returns_instance(x)
371366

372-
# TODO: `Inner` would probably be wrong here. Ideally, we would return a synthetic protocol
373-
# with a type of `int` for the `attr` attribute.
374-
reveal_type(inner) # revealed: Inner
367+
# TODO: Ideally, we would return a synthetic protocol with an `int` attribute.
368+
reveal_type(inner) # revealed: Unknown
375369

376370
# TODO: should ideally be `int`
377-
reveal_type(inner.attr) # revealed: T@returns_instance
371+
reveal_type(inner.attr) # revealed: Unknown
378372
```
379373

380374
## Generators

‎crates/ty_python_semantic/src/types.rs‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9576,6 +9576,51 @@ impl<'db> Type<'db> {
95769576
tcx: TypeContext<'db>,
95779577
visitor: &ApplyTypeMappingVisitor<'_, 'db>,
95789578
) -> Type<'db> {
9579+
if let TypeMapping::GeneralizeLocalTypes(scope) = type_mapping {
9580+
let is_local = |definition: Definition<'db>| {
9581+
definition.program_file(db) == scope.program_file(db)
9582+
&& semantic_index(db, scope.program_file(db))
9583+
.ancestor_scopes(definition.file_scope(db))
9584+
.any(|(id, _)| id == scope.file_scope_id(db))
9585+
};
9586+
9587+
let class = match self {
9588+
Type::ClassLiteral(class) => Some(class),
9589+
Type::GenericAlias(alias) => Some(alias.origin(db).into()),
9590+
Type::NominalInstance(instance) => Some(instance.class_literal(db, visitor.env)),
9591+
Type::SubclassOf(subclass) => match subclass.subclass_of() {
9592+
SubclassOfInner::Class(class) => Some(class.class_literal(db)),
9593+
_ => None,
9594+
},
9595+
_ => None,
9596+
};
9597+
if class
9598+
.and_then(|class| class.definition(db))
9599+
.is_some_and(is_local)
9600+
{
9601+
// See doc comment for `TypeMapping::GeneralizeLocalTypes`: we might want to
9602+
// create a synthetic protocol that represents the local class, instead of
9603+
// returning `Unknown`.
9604+
return Type::unknown();
9605+
}
9606+
9607+
if let Type::FunctionLiteral(function) = self
9608+
&& is_local(function.definition(db))
9609+
{
9610+
// Preserve the signature, but not the identity of a function object.
9611+
return visitor.visit(db, self, type_mapping, || {
9612+
Type::Callable(CallableType::new(
9613+
db,
9614+
function
9615+
.signature_with_inferred_return_type(db)
9616+
.into_owned(),
9617+
function.callable_type_kind(db),
9618+
))
9619+
.apply_type_mapping_impl(db, type_mapping, tcx, visitor)
9620+
});
9621+
}
9622+
}
9623+
95799624
// If we are binding `typing.Self`, and this type is what we are binding `Self` to, return
95809625
// early. This is not just an optimization, it also prevents us from infinitely expanding
95819626
// the type, if it's something that can contain a `Self` reference.
@@ -9916,6 +9961,7 @@ impl<'db> Type<'db> {
99169961
| TypeMapping::Materialize(_)
99179962
| TypeMapping::ReplaceParameterDefaults
99189963
| TypeMapping::EagerExpansion
9964+
| TypeMapping::GeneralizeLocalTypes(_)
99199965
| TypeMapping::RescopeReturnCallables(_)
99209966
| TypeMapping::Promote(PromotionMode::Off, _)
99219967
| TypeMapping::Promote(
@@ -9938,6 +9984,7 @@ impl<'db> Type<'db> {
99389984
| TypeMapping::Promote(..)
99399985
| TypeMapping::ReplaceParameterDefaults
99409986
| TypeMapping::EagerExpansion
9987+
| TypeMapping::GeneralizeLocalTypes(_)
99419988
| TypeMapping::RescopeReturnCallables(_) => self,
99429989
TypeMapping::Materialize(materialization_kind) => match materialization_kind {
99439990
MaterializationKind::Top => Type::object(),
@@ -11309,6 +11356,29 @@ pub enum TypeMapping<'a, 'db> {
1130911356

1131011357
/// Updates any `Callable` types in a function signature return type to be generic if possible.
1131111358
RescopeReturnCallables(&'a FxHashMap<CallableType<'db>, CallableType<'db>>),
11359+
11360+
/// This type mapping is used in return type inference. It "promotes" certain types to supertypes
11361+
/// or to dynamic types in order to prevent nested definitions from leaking into the inferred
11362+
/// return type.
11363+
///
11364+
/// For example, consider this unannotated decorator:
11365+
/// ```python
11366+
/// def register(event: str):
11367+
/// def wrapper(func):
11368+
/// # ...
11369+
/// return func
11370+
///
11371+
/// return wrapper
11372+
/// ```
11373+
///
11374+
/// We do not want to infer `def wrapper(func) -> Unknown` as the return type of `register`. This would
11375+
/// not just expose internal details, it would also incorrectly suggest that each call to `register`
11376+
/// returns the same function object. Instead, we generalize to an anonymous `(func) -> Unknown` callable.
11377+
///
11378+
/// For nested classes, we cannot easily do the same, so we return `Unknown` instead. In the future, to
11379+
/// improve this, we might want to synthesize a protocol that has the same attributes and methods as the
11380+
/// nested class.
11381+
GeneralizeLocalTypes(ScopeId<'db>),
1131211382
}
1131311383

1131411384
impl<'db> TypeMapping<'_, 'db> {
@@ -11372,6 +11442,7 @@ impl<'db> TypeMapping<'_, 'db> {
1137211442
| TypeMapping::Materialize(_)
1137311443
| TypeMapping::ReplaceParameterDefaults
1137411444
| TypeMapping::EagerExpansion
11445+
| TypeMapping::GeneralizeLocalTypes(_)
1137511446
| TypeMapping::RescopeReturnCallables(_) => Some(context),
1137611447
TypeMapping::BindSelf(binding) => {
1137711448
if binding.binding_context().is_some() {
@@ -11424,6 +11495,7 @@ impl<'db> TypeMapping<'_, 'db> {
1142411495
| TypeMapping::ReplaceSelf { .. }
1142511496
| TypeMapping::ReplaceParameterDefaults
1142611497
| TypeMapping::EagerExpansion
11498+
| TypeMapping::GeneralizeLocalTypes(_)
1142711499
| TypeMapping::RescopeReturnCallables(_) => self.clone(),
1142811500
}
1142911501
}

‎crates/ty_python_semantic/src/types/function.rs‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1794,7 +1794,14 @@ impl<'db> FunctionType<'db> {
17941794
let Some(class_definition) = index.class_definition_of_method(scope.file_scope_id(db))
17951795
else {
17961796
// For free functions, we simply use the type inferred from the function body.
1797-
return infer_complete_scope_types(db, scope).return_type();
1797+
return infer_complete_scope_types(db, scope)
1798+
.return_type()
1799+
.apply_type_mapping(
1800+
db,
1801+
&env,
1802+
&TypeMapping::GeneralizeLocalTypes(scope),
1803+
TypeContext::default(),
1804+
);
17981805
};
17991806
if self
18001807
.literal(db)
@@ -1833,7 +1840,14 @@ impl<'db> FunctionType<'db> {
18331840
return_type = UnionType::from_two_elements(db, &env, return_type, Type::unknown());
18341841
}
18351842

1836-
return_type.apply_optional_specialization(db, self.class_specialization(db))
1843+
return_type
1844+
.apply_type_mapping(
1845+
db,
1846+
&env,
1847+
&TypeMapping::GeneralizeLocalTypes(scope),
1848+
TypeContext::default(),
1849+
)
1850+
.apply_optional_specialization(db, self.class_specialization(db))
18371851
}
18381852

18391853
#[salsa::tracked(

‎crates/ty_python_semantic/src/types/known_instance.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -617,6 +617,7 @@ impl<'db> KnownInstanceType<'db> {
617617
| TypeMapping::Materialize(_)
618618
| TypeMapping::ReplaceParameterDefaults
619619
| TypeMapping::EagerExpansion
620+
| TypeMapping::GeneralizeLocalTypes(_)
620621
| TypeMapping::RescopeReturnCallables(_)
621622
| TypeMapping::ApplyRecursiveSubstitution(_) => Type::KnownInstance(self),
622623
},

‎crates/ty_python_semantic/src/types/typevar.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1482,6 +1482,7 @@ impl<'db> BoundTypeVarInstance<'db> {
14821482
| TypeMapping::ReplaceParameterDefaults
14831483
| TypeMapping::BindLegacyTypevars(_)
14841484
| TypeMapping::EagerExpansion
1485+
| TypeMapping::GeneralizeLocalTypes(_)
14851486
| TypeMapping::RescopeReturnCallables(_)
14861487
| TypeMapping::ApplyRecursiveSubstitution(_) => Type::TypeVar(self),
14871488
TypeMapping::Materialize(materialization_kind) => {

0 commit comments

Comments
 (0)