mirror of
https://github.com/astral-sh/ruff.git
synced 2026-05-06 08:56:57 -04:00
[ty] Selectively promote a union of homogeneous fixed-length tuples to a single variadic tuple. (#24705)
<!-- Thank you for contributing to Ruff/ty! To help us out with reviewing, please consider the following: - Does this pull request include a summary of the change? (See below.) - Does this pull request include a descriptive title? (Please prefix with `[ty]` for ty pull requests.) - Does this pull request include references to any relevant issues? - Does this PR follow our AI policy (https://github.com/astral-sh/.github/blob/main/AI_POLICY.md)? --> ## Summary This implements structural promotion of tuple size in the inferred type of a collection literal. The promotion only applies to a very specific circumstance: when tuple literals in an inferred collection element type produce a union of homogeneous fixed-length tuples of differing lengths, and only literal tuple sources have contributed to that type, then we widen that union to a single variadic tuple (e.g., `tuple[str] | tuple[str, str]` is widened to `tuple[str, ...]`). The result is that this scenario described in https://github.com/astral-sh/ty/issues/2620 succeeds: ```python languages = { "python": (".py", ".pyi"), "javascript": (".js", ".jsx", ".ts", ".tsx"), } # This no longer errors after this change, because the type of languages is `dict[str, tuple[str, ...]` rather than `dict[str, tuple[str, str]] | tuple[str, str, str, str]]` languages["ruby"] = (".rb",) ``` Closes https://github.com/astral-sh/ty/issues/2620. ### Approach - I created a new submodule that encapsulates the tuple size promotion policy. It exposes a `TupleSizePromotionConstraints` struct that we use during inference to record the scenarios in which we should **not** attempt to promote a tuple. If no such disqualifying scenarios are encountered, then tuple size promotion is attempted. The set of disqualifying scenarios is documented in new mdtests. - I think the policy for when to promote unions that involve empty tuples deserves particular scrutiny. Since empty tuples do not have an element type, they present a special case. The rule I've chosen is that empty tuples do not contribute to evidence of different tuple lengths. That means that a union containing an empty tuple must also contain other tuples of differing lengths to trigger promotion (i.e., `[(), (1,)]` remains `list[tuple[()] | tuple[int]]`, but `[(), (1,), (1,2)]` is promoted to `list[tuple[int, ...]]`. This is conservative and, I hope, useful for modeling situations in which the size of a tuple is specifically meant to be 0 or N. ## Test Plan Please see new and updated mdtests. <!-- How was it tested? -->
This commit is contained in:
Lérè
@@ -65,7 +65,7 @@ reveal_type(x) # revealed: dict[int, (_: int) -> int]
|
||||
## Mixed dict
|
||||
|
||||
```py
|
||||
# revealed: dict[str, int | tuple[int, int] | tuple[int, int, int]]
|
||||
# revealed: dict[str, int | tuple[int, ...]]
|
||||
reveal_type({"a": 1, "b": (1, 2), "c": (1, 2, 3)})
|
||||
```
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ reveal_type(x[0].__name__) # revealed: str
|
||||
## Mixed list
|
||||
|
||||
```py
|
||||
# revealed: list[int | tuple[int, int] | tuple[int, int, int]]
|
||||
# revealed: list[int | tuple[int, ...]]
|
||||
reveal_type([1, (1, 2), (1, 2, 3)])
|
||||
```
|
||||
|
||||
|
||||
@@ -28,7 +28,7 @@ reveal_type(x) # revealed: set[(_: int) -> int]
|
||||
## Mixed set
|
||||
|
||||
```py
|
||||
# revealed: set[int | tuple[int, int] | tuple[int, int, int]]
|
||||
# revealed: set[int | tuple[int, ...]]
|
||||
reveal_type({1, (1, 2), (1, 2, 3)})
|
||||
```
|
||||
|
||||
|
||||
@@ -90,6 +90,209 @@ reveal_type((1, 2, 3)) # revealed: tuple[Literal[1], Literal[2], Literal[3]]
|
||||
reveal_type(frozenset((1, 2, 3))) # revealed: frozenset[Literal[1, 2, 3]]
|
||||
```
|
||||
|
||||
## Unions of homogeneous, fixed-length tuples can be promoted to a single variadic tuple
|
||||
|
||||
This type of promotion applies specifically when a collection literal contains at least two tuple
|
||||
literals that share an element type but that differ in length. The inferred type of those literals
|
||||
is a union of homogeneous, fixed-length tuples, which is subsequently promoted to a single, variadic
|
||||
tuple. The type of any non-tuple elements in the collection literal is preserved in the final
|
||||
inferred type.
|
||||
|
||||
```py
|
||||
reveal_type([(1, 2), (3, 4, 5)]) # revealed: list[tuple[int, ...]]
|
||||
reveal_type({".py": (".py", ".pyi"), ".js": (".js", ".jsx", ".ts", ".tsx")}) # revealed: dict[str, tuple[str, ...]]
|
||||
reveal_type({(1, 2), (3, 4, 5)}) # revealed: set[tuple[int, ...]]
|
||||
reveal_type([0, (1, 2), "a", (3, 4, 5)]) # revealed: list[int | str | tuple[int, ...]]
|
||||
```
|
||||
|
||||
We only widen unions of fixed-length tuples. A standalone tuple type retains its fixed length.
|
||||
|
||||
```py
|
||||
def promote[T](x: T) -> list[T]:
|
||||
return [x]
|
||||
|
||||
reveal_type([()]) # revealed: list[tuple[()]]
|
||||
reveal_type((1, 2)) # revealed: tuple[Literal[1], Literal[2]]
|
||||
reveal_type(promote((1, 2))) # revealed: list[tuple[int, int]]
|
||||
```
|
||||
|
||||
Tuple literals of the same length also keep their fixed length. For example, a declared collection
|
||||
representing coordinate pairs does not later accept a coordinate triple.
|
||||
|
||||
```py
|
||||
coordinates = {
|
||||
"home": (0, 0),
|
||||
"palm-tree": (10, 8),
|
||||
}
|
||||
reveal_type(coordinates) # revealed: dict[str, tuple[int, int]]
|
||||
coordinates["treasure"] = (5, 6, -10) # error: [invalid-assignment]
|
||||
```
|
||||
|
||||
Heterogeneous tuples are not widened.
|
||||
|
||||
```py
|
||||
reveal_type([(1, "a"), (2, "b")]) # revealed: list[tuple[int, str]]
|
||||
reveal_type([(1, 2), ("a", "b", "c")]) # revealed: list[tuple[int, int] | tuple[str, str, str]]
|
||||
```
|
||||
|
||||
Normal union simplification can still generalize heterogeneous tuples of the same length, but the
|
||||
result should not widen to a variadic tuple.
|
||||
|
||||
```py
|
||||
mixed_tuples = [
|
||||
(1, "a"),
|
||||
(object(), object()),
|
||||
(object(), object(), object()),
|
||||
]
|
||||
reveal_type(mixed_tuples) # revealed: list[tuple[object, object] | tuple[object, object, object]]
|
||||
```
|
||||
|
||||
Empty tuples are treated specially: they do not count towards the minimum of two tuples that differ
|
||||
in length. This allows us to preserve accurate types for collections that model exactly two possible
|
||||
tuple shapes (empty or length N).
|
||||
|
||||
```py
|
||||
reveal_type([(), (1,)]) # revealed: list[tuple[()] | tuple[int]]
|
||||
reveal_type([(), (1,), (2,)]) # revealed: list[tuple[()] | tuple[int]]
|
||||
```
|
||||
|
||||
However, an empty tuple can be subsumed by a widened tuple type when enough evidence exists
|
||||
independently of the empty tuple.
|
||||
|
||||
```py
|
||||
reveal_type([(), (1, 2), (1, 2, 3)]) # revealed: list[tuple[int, ...]]
|
||||
```
|
||||
|
||||
A union of tuples is not widened when it is inferred into a collection (i.e., only unions that came
|
||||
from direct tuple literals in the collection are widened).
|
||||
|
||||
```py
|
||||
def get_padding() -> int | tuple[int] | tuple[int, int]:
|
||||
return (0, 1)
|
||||
|
||||
reveal_type([get_padding()]) # revealed: list[int | tuple[int] | tuple[int, int]]
|
||||
```
|
||||
|
||||
No promotion occurs in a collection that mixes literal and non-literal tuples:
|
||||
|
||||
```py
|
||||
def get_segment() -> tuple[int] | tuple[int, int, int, int]:
|
||||
return (0,)
|
||||
|
||||
def get_segments() -> list[tuple[int, int]]:
|
||||
return [(0, 1)]
|
||||
|
||||
def get_segments_by_name() -> dict[str, tuple[int, int]]:
|
||||
return {"origin": (0, 1)}
|
||||
|
||||
segments = [get_segment(), (1, 2), (3, 4, 5)]
|
||||
reveal_type(segments) # revealed: list[tuple[int] | tuple[int, int, int, int] | tuple[int, int] | tuple[int, int, int]]
|
||||
segments.append((6, 7, 8, 9, 10)) # error: [invalid-argument-type]
|
||||
|
||||
starred_segments = [*get_segments(), (1, 2), (3, 4, 5)]
|
||||
reveal_type(starred_segments) # revealed: list[tuple[int, int] | tuple[int, int, int]]
|
||||
starred_segments.append((6, 7, 8, 9)) # error: [invalid-argument-type]
|
||||
|
||||
mapping_segments = {**get_segments_by_name(), "start": (1, 2), "end": (3, 4, 5)}
|
||||
reveal_type(mapping_segments) # revealed: dict[str, tuple[int, int] | tuple[int, int, int]]
|
||||
mapping_segments["bad"] = (6, 7, 8, 9) # error: [invalid-assignment]
|
||||
```
|
||||
|
||||
This also applies when the non-literal tuple type is hidden behind a type alias or a type variable,
|
||||
or when it is subsumed by one of the literal tuple types while building the union.
|
||||
|
||||
```toml
|
||||
[environment]
|
||||
python-version = "3.12"
|
||||
```
|
||||
|
||||
```py
|
||||
from typing import NewType, TypeVar
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
type Segment = tuple[int, int]
|
||||
NewTypeSegment = NewType("NewTypeSegment", tuple[int, int])
|
||||
BoundSegment = TypeVar("BoundSegment", bound=tuple[int, int])
|
||||
ConstrainedSegment = TypeVar("ConstrainedSegment", tuple[int, int], tuple[int, int, int])
|
||||
|
||||
class P: ...
|
||||
|
||||
def is_p(x: object) -> TypeIs[P]:
|
||||
return True
|
||||
|
||||
def get_aliased_segment() -> Segment:
|
||||
return (0, 1)
|
||||
|
||||
aliased_segments = [get_aliased_segment(), (1, 2), (3, 4, 5)]
|
||||
reveal_type(aliased_segments) # revealed: list[tuple[int, int] | tuple[int, int, int]]
|
||||
aliased_segments.append((6, 7, 8, 9)) # error: [invalid-argument-type]
|
||||
|
||||
def get_newtype_segment() -> NewTypeSegment:
|
||||
return NewTypeSegment((0, 1))
|
||||
|
||||
newtype_segments = [get_newtype_segment(), (1, 2), (3, 4, 5)]
|
||||
reveal_type(newtype_segments) # revealed: list[tuple[int, int] | tuple[int, int, int]]
|
||||
newtype_segments.append((6, 7, 8, 9)) # error: [invalid-argument-type]
|
||||
|
||||
def check_bound_typevar_segment(segment: BoundSegment) -> None:
|
||||
bound_typevar_segments = [segment, (1, 2), (3, 4, 5)]
|
||||
reveal_type(bound_typevar_segments) # revealed: list[tuple[int, int] | tuple[int, int, int]]
|
||||
bound_typevar_segments.append((6, 7, 8, 9)) # error: [invalid-argument-type]
|
||||
|
||||
def check_constrained_typevar_segment(segment: ConstrainedSegment) -> None:
|
||||
constrained_typevar_segments = [segment, (1, 2), (3, 4, 5)]
|
||||
# revealed: list[ConstrainedSegment@check_constrained_typevar_segment | tuple[int, int] | tuple[int, int, int]]
|
||||
reveal_type(constrained_typevar_segments)
|
||||
constrained_typevar_segments.append((6, 7, 8, 9)) # error: [invalid-argument-type]
|
||||
|
||||
def get_subsumed_segment() -> tuple[bool, bool]:
|
||||
return (True, False)
|
||||
|
||||
subsumed_segments = [get_subsumed_segment(), (1, 2), (3, 4, 5)]
|
||||
reveal_type(subsumed_segments) # revealed: list[tuple[int, int] | tuple[int, int, int]]
|
||||
subsumed_segments.append((6, 7, 8, 9)) # error: [invalid-argument-type]
|
||||
|
||||
def get_short_subsumed_segment() -> tuple[bool]:
|
||||
return (True,)
|
||||
|
||||
short_subsumed_segments = [get_short_subsumed_segment(), (1, 2), (3, 4, 5)]
|
||||
reveal_type(short_subsumed_segments) # revealed: list[tuple[bool] | tuple[int, int] | tuple[int, int, int]]
|
||||
short_subsumed_segments.append((6, 7, 8, 9)) # error: [invalid-argument-type]
|
||||
|
||||
def get_heterogeneous_subsumed_segment() -> tuple[bool, int]:
|
||||
return (True, 0)
|
||||
|
||||
heterogeneous_subsumed_segments = [get_heterogeneous_subsumed_segment(), (1, 2), (3, 4, 5)]
|
||||
reveal_type(heterogeneous_subsumed_segments) # revealed: list[tuple[int, int] | tuple[int, int, int]]
|
||||
heterogeneous_subsumed_segments.append((6, 7, 8, 9)) # error: [invalid-argument-type]
|
||||
|
||||
def check_intersection_segment(segment: tuple[int, int]) -> None:
|
||||
if is_p(segment):
|
||||
reveal_type(segment) # revealed: tuple[int, int] & P
|
||||
intersection_segments = [segment, (1, 2), (3, 4, 5)]
|
||||
reveal_type(intersection_segments) # revealed: list[tuple[int, int] | tuple[int, int, int]]
|
||||
intersection_segments.append((6, 7, 8, 9)) # error: [invalid-argument-type]
|
||||
```
|
||||
|
||||
No promotion occurs when a covariant collection type context provides a fixed-length tuple.
|
||||
|
||||
```py
|
||||
from typing import Sequence
|
||||
|
||||
segments: Sequence[tuple[int, int] | tuple[int, int, int]] = [(1, 2), (3, 4, 5)]
|
||||
reveal_type(segments) # revealed: list[tuple[int, int] | tuple[int, int, int]]
|
||||
```
|
||||
|
||||
Promotion still occurs when a covariant collection type context provides a gradual element type.
|
||||
|
||||
```py
|
||||
from typing import Any
|
||||
from collections.abc import Mapping
|
||||
|
||||
segments: Mapping[str, Any] = {"start": (1, 2), "end": (3, 4, 5)}
|
||||
reveal_type(segments) # revealed: dict[str, tuple[int, ...]]
|
||||
```
|
||||
|
||||
## Invariant and contravariant return types are promoted
|
||||
|
||||
We promote in non-covariant position in the return type of a generic function, or constructor of a
|
||||
|
||||
@@ -84,6 +84,7 @@ use crate::types::set_theoretic::RecursivelyDefined;
|
||||
use crate::types::signatures::CallableSignature;
|
||||
use crate::types::special_form::TypeQualifier;
|
||||
use crate::types::subclass_of::SubclassOfInner;
|
||||
use crate::types::tuple::promotion::TupleSizePromotionConstraints;
|
||||
use crate::types::tuple::{Tuple, TupleLength, TupleSpecBuilder, TupleType};
|
||||
use crate::types::type_alias::{ManualPEP695TypeAliasType, PEP695TypeAliasType};
|
||||
use crate::types::typevar::{BoundTypeVarIdentity, TypeVarConstraints, TypeVarIdentity};
|
||||
@@ -6165,6 +6166,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
|
||||
// Create a set of constraints to infer a precise type for `T`.
|
||||
let mut builder = SpecializationBuilder::new(self.db(), &constraints, inferable);
|
||||
|
||||
let mut tuple_size_promotion_constraints = TupleSizePromotionConstraints::default();
|
||||
|
||||
for elt_ty in elt_tys.clone() {
|
||||
let elt_ty_identity = elt_ty.identity(self.db());
|
||||
let elt_tcx = elt_tcx_constraints
|
||||
@@ -6175,6 +6178,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
|
||||
.get(&elt_ty_identity)
|
||||
.copied();
|
||||
|
||||
if elt_tcx.is_some_and(|elt_tcx| !elt_tcx.is_dynamic()) {
|
||||
// Record type annotations that provide concrete shape information in order to
|
||||
// disqualify this typevar from tuple size promotion.
|
||||
tuple_size_promotion_constraints.record_declared_type(elt_ty_identity);
|
||||
}
|
||||
|
||||
// Avoid unnecessarily widening the return type based on a covariant
|
||||
// type parameter from the type context.
|
||||
//
|
||||
@@ -6221,6 +6230,17 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
|
||||
|
||||
let mut elt_tys = elt_tys.clone();
|
||||
if let Some((key_ty, value_ty)) = elt_tys.next_tuple() {
|
||||
tuple_size_promotion_constraints.record_unpromotable_type(
|
||||
self.db(),
|
||||
key_ty.identity(self.db()),
|
||||
unpacked_key_ty.promote(self.db()),
|
||||
);
|
||||
tuple_size_promotion_constraints.record_unpromotable_type(
|
||||
self.db(),
|
||||
value_ty.identity(self.db()),
|
||||
unpacked_value_ty.promote(self.db()),
|
||||
);
|
||||
|
||||
builder.infer(Type::TypeVar(key_ty), unpacked_key_ty).ok()?;
|
||||
|
||||
builder
|
||||
@@ -6269,17 +6289,23 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
|
||||
// which the constraint solver struggles with.
|
||||
let inferred_elt_ty = inferred_elt_ty.promote(self.db());
|
||||
|
||||
let inferred_type_for_typevar = if elt.is_starred_expr() {
|
||||
inferred_elt_ty
|
||||
.iterate(self.db())
|
||||
.homogeneous_element_type(self.db())
|
||||
} else {
|
||||
inferred_elt_ty
|
||||
};
|
||||
|
||||
tuple_size_promotion_constraints.record_inferred_expression_type(
|
||||
self.db(),
|
||||
elt_ty_identity,
|
||||
elt,
|
||||
inferred_type_for_typevar,
|
||||
);
|
||||
|
||||
builder
|
||||
.infer(
|
||||
Type::TypeVar(elt_ty),
|
||||
if elt.is_starred_expr() {
|
||||
inferred_elt_ty
|
||||
.iterate(self.db())
|
||||
.homogeneous_element_type(self.db())
|
||||
} else {
|
||||
inferred_elt_ty
|
||||
},
|
||||
)
|
||||
.infer(Type::TypeVar(elt_ty), inferred_type_for_typevar)
|
||||
.ok()?;
|
||||
}
|
||||
}
|
||||
@@ -6287,14 +6313,27 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
|
||||
let class_type = collection_alias
|
||||
.origin(self.db())
|
||||
.apply_specialization(self.db(), |_| {
|
||||
builder.build_with(generic_context, |_, bounds| {
|
||||
builder.build_with(generic_context, |current_typevar, bounds| {
|
||||
let (lower, _upper) = bounds?;
|
||||
// Promote singleton types to `T | Unknown` in inferred type parameters,
|
||||
// so that e.g. `[None]` is inferred as `list[None | Unknown]`.
|
||||
if elt_tcx_constraints.is_empty() {
|
||||
return Some(lower.promote_singletons_recursively(self.db()));
|
||||
}
|
||||
None
|
||||
|
||||
let lower = if tuple_size_promotion_constraints
|
||||
.allow(current_typevar.identity(self.db()))
|
||||
{
|
||||
lower.promote_tuple_size_in_union(self.db())
|
||||
} else {
|
||||
lower
|
||||
};
|
||||
|
||||
let lower = if elt_tcx_constraints.is_empty() {
|
||||
lower
|
||||
// Promote singleton types to `T | Unknown` in inferred type parameters,
|
||||
// so that e.g. `[None]` is inferred as `list[None | Unknown]`.
|
||||
.promote_singletons_recursively(self.db())
|
||||
} else {
|
||||
lower
|
||||
};
|
||||
|
||||
Some(lower)
|
||||
})
|
||||
});
|
||||
|
||||
|
||||
@@ -35,6 +35,8 @@ use crate::{Db, FxOrderSet, Program};
|
||||
use ty_python_core::Truthiness;
|
||||
use ty_python_core::definition::Definition;
|
||||
|
||||
pub(crate) mod promotion;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum TupleLength {
|
||||
Fixed(usize),
|
||||
|
||||
@@ -0,0 +1,210 @@
|
||||
use rustc_hash::FxHashSet;
|
||||
|
||||
use ruff_python_ast as ast;
|
||||
|
||||
use crate::Db;
|
||||
use crate::types::tuple::TupleSpec;
|
||||
use crate::types::typevar::BoundTypeVarIdentity;
|
||||
use crate::types::visitor::any_over_type;
|
||||
use crate::types::{Type, UnionBuilder};
|
||||
|
||||
/// Tracks the typevars of a collection to which tuple size promotion should **not** apply.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct TupleSizePromotionConstraints<'db> {
|
||||
blocked_typevars: FxHashSet<BoundTypeVarIdentity<'db>>,
|
||||
}
|
||||
|
||||
impl<'db> TupleSizePromotionConstraints<'db> {
|
||||
/// Records that a typevar has a declared type. This makes it ineligible for tuple size promotion.
|
||||
pub(crate) fn record_declared_type(&mut self, typevar_identity: BoundTypeVarIdentity<'db>) {
|
||||
self.blocked_typevars.insert(typevar_identity);
|
||||
}
|
||||
|
||||
/// Records whether an inferred collection element blocks tuple size promotion for the typevar.
|
||||
pub(crate) fn record_inferred_expression_type(
|
||||
&mut self,
|
||||
db: &'db dyn Db,
|
||||
typevar_identity: BoundTypeVarIdentity<'db>,
|
||||
expression: &ast::Expr,
|
||||
ty: Type<'db>,
|
||||
) {
|
||||
if !Self::is_promotable_tuple_literal(db, expression, ty) {
|
||||
self.record_unpromotable_type(db, typevar_identity, ty);
|
||||
}
|
||||
}
|
||||
|
||||
/// Records that a typevar is ineligible for tuple size promotion if the given type contains
|
||||
/// a tuple type.
|
||||
pub(crate) fn record_unpromotable_type(
|
||||
&mut self,
|
||||
db: &'db dyn Db,
|
||||
typevar_identity: BoundTypeVarIdentity<'db>,
|
||||
ty: Type<'db>,
|
||||
) {
|
||||
if any_over_type(db, ty, true, |ty| ty.tuple_instance_spec(db).is_some()) {
|
||||
self.blocked_typevars.insert(typevar_identity);
|
||||
}
|
||||
}
|
||||
|
||||
/// Reports whether or not tuple size promotion is allowed for the given typevar in light
|
||||
/// of the constraints recorded on this object.
|
||||
pub(crate) fn allow(&self, typevar_identity: BoundTypeVarIdentity<'db>) -> bool {
|
||||
!self.blocked_typevars.contains(&typevar_identity)
|
||||
}
|
||||
|
||||
/// Returns true if the given expression is either a non-starred homogeneous tuple literal or the
|
||||
/// empty tuple (and hence is eligible for tuple size promotion).
|
||||
fn is_promotable_tuple_literal(db: &'db dyn Db, expression: &ast::Expr, ty: Type<'db>) -> bool {
|
||||
matches!(expression, ast::Expr::Tuple(tuple) if !tuple.iter().any(ast::Expr::is_starred_expr))
|
||||
&& TupleSizePromotionCandidate::from_type(db, ty).is_some()
|
||||
}
|
||||
}
|
||||
|
||||
/// Represents a single tuple literal whose type in the inferred collection type might be widened.
|
||||
enum TupleSizePromotionCandidate<'db> {
|
||||
Empty,
|
||||
Homogeneous {
|
||||
element_type: Type<'db>,
|
||||
length: usize,
|
||||
},
|
||||
}
|
||||
|
||||
impl<'db> TupleSizePromotionCandidate<'db> {
|
||||
/// Returns an eligible candidate if the given type represents one (i.e., it is a
|
||||
/// fixed-length homogeneous tuple or the empty tuple).
|
||||
fn from_type(db: &'db dyn Db, ty: Type<'db>) -> Option<Self> {
|
||||
let tuple_spec = ty.exact_tuple_instance_spec(db)?;
|
||||
let TupleSpec::Fixed(tuple) = tuple_spec.as_ref() else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut elements = tuple.iter_all_elements();
|
||||
let Some(element_type) = elements.next() else {
|
||||
return Some(Self::Empty);
|
||||
};
|
||||
|
||||
elements
|
||||
.all(|element| element.is_equivalent_to(db, element_type))
|
||||
.then_some(Self::Homogeneous {
|
||||
element_type,
|
||||
length: tuple.len(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Represents a group of tuple types extracted from a larger union. The types in this group may
|
||||
/// be widened in the final inferred type for the collection literal.
|
||||
struct HomogeneousTupleUnionGroup<'db> {
|
||||
element_type: Type<'db>,
|
||||
original_tuple_types: Vec<Type<'db>>,
|
||||
first_length: usize,
|
||||
has_multiple_lengths: bool,
|
||||
}
|
||||
|
||||
impl<'db> HomogeneousTupleUnionGroup<'db> {
|
||||
fn new(element_type: Type<'db>, original_tuple_type: Type<'db>, length: usize) -> Self {
|
||||
Self {
|
||||
element_type,
|
||||
original_tuple_types: vec![original_tuple_type],
|
||||
first_length: length,
|
||||
has_multiple_lengths: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Adds a tuple to this homogeneous union group.
|
||||
fn add(&mut self, original_tuple_type: Type<'db>, length: usize) {
|
||||
self.has_multiple_lengths |= length != self.first_length;
|
||||
self.original_tuple_types.push(original_tuple_type);
|
||||
}
|
||||
}
|
||||
|
||||
/// Partitions a union into two sets prior to rebuilding it: one for elements that are not
|
||||
/// candidates for tuple size promotion, and another for groups of homogeneous tuple elements that are.
|
||||
fn partition_tuple_union_elements<'db>(
|
||||
db: &'db dyn Db,
|
||||
elements: impl IntoIterator<Item = Type<'db>>,
|
||||
) -> (Vec<Type<'db>>, Vec<HomogeneousTupleUnionGroup<'db>>) {
|
||||
let mut other_union_elements = Vec::new();
|
||||
let mut tuple_groups: Vec<HomogeneousTupleUnionGroup<'db>> = Vec::new();
|
||||
|
||||
for element in elements {
|
||||
match TupleSizePromotionCandidate::from_type(db, element) {
|
||||
Some(TupleSizePromotionCandidate::Homogeneous {
|
||||
element_type,
|
||||
length,
|
||||
}) => {
|
||||
if let Some(group) = tuple_groups
|
||||
.iter_mut()
|
||||
.find(|group| group.element_type.is_equivalent_to(db, element_type))
|
||||
{
|
||||
group.add(element, length);
|
||||
} else {
|
||||
tuple_groups.push(HomogeneousTupleUnionGroup::new(
|
||||
element_type,
|
||||
element,
|
||||
length,
|
||||
));
|
||||
}
|
||||
}
|
||||
Some(TupleSizePromotionCandidate::Empty) | None => other_union_elements.push(element),
|
||||
}
|
||||
}
|
||||
|
||||
(other_union_elements, tuple_groups)
|
||||
}
|
||||
|
||||
impl<'db> Type<'db> {
|
||||
/// Within a larger union, promotes every group of homogeneous, fixed-length tuples of differing
|
||||
/// lengths to a single variadic tuple.
|
||||
///
|
||||
/// This deliberately only applies to unions; a standalone tuple keeps its shape.
|
||||
///
|
||||
/// The caller is responsible for checking that every tuple source that contributes to this
|
||||
/// union is eligible for promotion (see [`TupleSizePromotionConstraints`]).
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// In the code below, we promote `dict[str, tuple[str, str] | tuple[str, str, str, str]]`
|
||||
/// to `dict[str, tuple[str, ...]]`:
|
||||
///
|
||||
/// ```python
|
||||
/// languages = {
|
||||
/// "python": (".py", ".pyi"),
|
||||
/// "javascript": (".js", ".jsx", ".ts", ".tsx"),
|
||||
/// }
|
||||
/// reveal_type(languages) # revealed: dict[str, tuple[str, ...]]
|
||||
/// ```
|
||||
///
|
||||
pub(crate) fn promote_tuple_size_in_union(self, db: &'db dyn Db) -> Type<'db> {
|
||||
let Type::Union(union) = self else {
|
||||
return self;
|
||||
};
|
||||
|
||||
let (other_union_elements, tuple_groups) =
|
||||
partition_tuple_union_elements(db, union.elements(db).iter().copied());
|
||||
|
||||
if !tuple_groups.iter().any(|group| group.has_multiple_lengths) {
|
||||
return self;
|
||||
}
|
||||
|
||||
let mut builder = UnionBuilder::new(db)
|
||||
.unpack_aliases(false)
|
||||
.recursively_defined(union.recursively_defined(db));
|
||||
|
||||
for element in other_union_elements {
|
||||
builder = builder.add(element);
|
||||
}
|
||||
|
||||
for group in tuple_groups {
|
||||
if group.has_multiple_lengths {
|
||||
builder = builder.add(Type::homogeneous_tuple(db, group.element_type));
|
||||
} else {
|
||||
for element in group.original_tuple_types {
|
||||
builder = builder.add(element);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
builder.build()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user