diff --git a/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md b/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md index 412243139f..34347ea1aa 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md @@ -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 +``` diff --git a/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md b/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md index eb05373917..2645221f01 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md @@ -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 +``` diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index b0d33ee097..378583d9ae 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -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(_) => {}