[ty] Add precise inference for unpacking a TypeVar if the TypeVar has an upper bound with a precise tuple spec (#19985)

This commit is contained in:
Alex Waygood 2025-08-19 22:11:30 +01:00 committed by GitHub
parent c82e255ca8
commit 662d18bd05
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 94 additions and 2 deletions

View file

@ -455,3 +455,45 @@ def overloaded_outer(t: T | None = None) -> None:
if t is not None:
inner(t)
```
## Unpacking a TypeVar
We can infer precise heterogeneous types from the result of an unpacking operation applied to a type
variable if the type variable's upper bound is a type with a precise tuple spec:
```py
from dataclasses import dataclass
from typing import NamedTuple, Final, TypeVar, Generic
T = TypeVar("T", bound=tuple[int, str])
def f(x: T) -> T:
a, b = x
reveal_type(a) # revealed: int
reveal_type(b) # revealed: str
return x
@dataclass
class Team(Generic[T]):
employees: list[T]
def x(team: Team[T]) -> Team[T]:
age, name = team.employees[0]
reveal_type(age) # revealed: int
reveal_type(name) # revealed: str
return team
class Age(int): ...
class Name(str): ...
class Employee(NamedTuple):
age: Age
name: Name
EMPLOYEES: Final = (Employee(name=Name("alice"), age=Age(42)),)
team = Team(employees=list(EMPLOYEES))
reveal_type(team.employees) # revealed: list[Employee]
age, name = team.employees[0]
reveal_type(age) # revealed: Age
reveal_type(name) # revealed: Name
```

View file

@ -454,3 +454,43 @@ def overloaded_outer[T](t: T | None = None) -> None:
if t is not None:
inner(t)
```
## Unpacking a TypeVar
We can infer precise heterogeneous types from the result of an unpacking operation applied to a
TypeVar if the TypeVar's upper bound is a type with a precise tuple spec:
```py
from dataclasses import dataclass
from typing import NamedTuple, Final
def f[T: tuple[int, str]](x: T) -> T:
a, b = x
reveal_type(a) # revealed: int
reveal_type(b) # revealed: str
return x
@dataclass
class Team[T: tuple[int, str]]:
employees: list[T]
def x[T: tuple[int, str]](team: Team[T]) -> Team[T]:
age, name = team.employees[0]
reveal_type(age) # revealed: int
reveal_type(name) # revealed: str
return team
class Age(int): ...
class Name(str): ...
class Employee(NamedTuple):
age: Age
name: Name
EMPLOYEES: Final = (Employee(name=Name("alice"), age=Age(42)),)
team = Team(employees=list(EMPLOYEES))
reveal_type(team.employees) # revealed: list[Employee]
age, name = team.employees[0]
reveal_type(age) # revealed: Age
reveal_type(name) # revealed: Name
```

View file

@ -4870,6 +4870,18 @@ impl<'db> Type<'db> {
Type::TypeAlias(alias) => {
return alias.value_type(db).try_iterate_with_mode(db, mode);
}
Type::NonInferableTypeVar(tvar) => match tvar.typevar(db).bound_or_constraints(db) {
Some(TypeVarBoundOrConstraints::UpperBound(bound)) => {
return bound.try_iterate_with_mode(db, mode);
}
// TODO: could we create a "union of tuple specs"...?
// (Same question applies to the `Type::Union()` branch lower down)
Some(TypeVarBoundOrConstraints::Constraints(_)) | None => {}
},
Type::TypeVar(_) => unreachable!(
"should not be able to iterate over type variable {} in inferable position",
self.display(db)
),
Type::Dynamic(_)
| Type::FunctionLiteral(_)
| Type::GenericAlias(_)
@ -4895,8 +4907,6 @@ impl<'db> Type<'db> {
| Type::EnumLiteral(_)
| Type::LiteralString
| Type::BytesLiteral(_)
| Type::TypeVar(_)
| Type::NonInferableTypeVar(_)
| Type::BoundSuper(_)
| Type::TypeIs(_)
| Type::TypedDict(_) => {}