Promote unions of homogenously-typed tuples of differing lengths to a single, variable size tuple

Author: lerebear

crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md

@@ -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)})
 ```
 

crates/ty_python_semantic/resources/mdtest/literal/collections/list.md

@@ -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)])
 ```
 

crates/ty_python_semantic/resources/mdtest/literal/collections/set.md

@@ -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)})
 ```
 

crates/ty_python_semantic/resources/mdtest/promotion.md

@@ -90,6 +90,96 @@ 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 differently-sized but homogenously-typed tuples are promoted to a single, variably-sized tuple when they appear in an invariant position
+
+When a union of homogeneously-typed tuples of different lengths appears in a promotable position,
+the union is replaced with a single tuple of variable length after promotion:
+
+```py
+from typing import Sequence
+
+class Invariant[T]:
+    x: T
+
+    def __init__(self, value: T): ...
+
+def promote[T](x: T) -> list[T]:
+    return [x]
+
+def make_invariant[T](x: T) -> Invariant[T]:
+    return Invariant(x)
+
+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, ...]]
+
+languages = {
+    "python": (".py", ".pyi"),
+    "javascript": (".js", ".jsx", ".ts", ".tsx"),
+}
+reveal_type(languages)  # revealed: dict[str, tuple[str, ...]]
+languages["ruby"] = (".rb",)
+
+covariant_segments: Sequence[tuple[int, ...]] = [(1, 2), (3, 4, 5)]
+reveal_type(covariant_segments)  # revealed: list[tuple[int, ...]]
+```
+
+Tuple-size promotion only applies when a promotable position sees differently-sized homogeneous
+tuples. Plain tuple literals stay fixed and literal-precise, same-sized homogeneous tuples stay
+fixed, heterogeneous tuples stay fixed, and empty tuples are not promoted on their own:
+
+```py
+reveal_type((1, 2))  # revealed: tuple[Literal[1], Literal[2]]
+reveal_type(promote((1, 2)))  # revealed: list[tuple[int, int]]
+reveal_type(make_invariant((1, 2)))  # revealed: Invariant[tuple[int, int]]
+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]]
+reveal_type([()])  # revealed: list[tuple[()]]
+
+places_of_interest = {
+    "home": (0, 0),
+    "palm-tree": (10, 8),
+}
+reveal_type(places_of_interest)  # revealed: dict[str, tuple[int, int]]
+places_of_interest["treasure"] = (5, 6, -10)  # error: [invalid-assignment]
+```
+
+Tuple-size promotion is based on tuple literal elements in the collection. Explicit finite tuple
+unions are not widened just because they are inferred into a collection:
+
+```py
+def get_padding() -> int | tuple[int] | tuple[int, int] | tuple[int, int, int, int]:
+    return (0, 1)
+
+def get_segment() -> tuple[int] | tuple[int, int, int, int] | tuple[int, int, int, int, int]:
+    return (0,)
+
+reveal_type([get_padding()])  # revealed: list[int | tuple[int] | tuple[int, int] | tuple[int, int, int, int]]
+reveal_type([get_segment()])  # revealed: list[tuple[int] | tuple[int, int, int, int] | tuple[int, int, int, int, int]]
+
+mixed_segments = [get_segment(), (1, 2), (3, 4, 5)]
+mixed_segments.append((6, 7, 8, 9, 10))
+mixed_segments.append((6, 7, 8, 9, 10, 11))  # error: [invalid-argument-type]
+
+def get_overlapping_segment() -> tuple[int, int] | tuple[int, int, int]:
+    return (0, 1)
+
+overlapping_segments = [get_overlapping_segment(), (1, 2), (3, 4, 5)]
+reveal_type(overlapping_segments)  # revealed: list[tuple[int, int] | tuple[int, int, int]]
+overlapping_segments.append((6, 7, 8, 9))  # error: [invalid-argument-type]
+
+def accepts_padding(padding: int | tuple[int] | tuple[int, int] | tuple[int, int, int, int]) -> None: ...
+
+class UsesPadding:
+    def __init__(self, padding: int | tuple[int] | tuple[int, int] | tuple[int, int, int, int]):
+        self.padding = padding
+
+    def check(self) -> None:
+        accepts_padding(self.padding)
+
+UsesPadding(get_padding()).check()
+```
+
 ## 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
@@ -254,6 +344,9 @@ reveal_type(x13)  # revealed: list[tuple[Literal[1], Literal[2], Literal[3]]]
 x14: list[tuple[int, str, int]] = [(1, "2", 3), (4, "5", 6)]
 reveal_type(x14)  # revealed: list[tuple[int, str, int]]
 
+x14a: list[tuple[int, int]] = [(1, 2), (3, 4)]
+reveal_type(x14a)  # revealed: list[tuple[int, int]]
+
 x15: list[tuple[Literal[1], ...]] = [(1, 1, 1)]
 reveal_type(x15)  # revealed: list[tuple[Literal[1], ...]]
 
@@ -543,6 +636,7 @@ any errors from clearly incorrect code like this:
 `module1.py`:
 
 ```py
+
 ```
 
 `main.py`:

crates/ty_python_semantic/src/types.rs

@@ -1,7 +1,7 @@
 use compact_str::ToCompactString;
 use itertools::Itertools;
 use ruff_diagnostics::{Edit, Fix};
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use std::borrow::Cow;
 use std::cell::OnceCell;
@@ -862,6 +862,84 @@ fn recursive_type_normalize_type_guard_like<'db, T: TypeGuardLike<'db>>(
     };
     Some(guard.with_type(db, ty))
 }
+// Represents the members of a union that are all homogeneously-typed, fixed-length tuples and that
+// share the same element type.
+//
+// EXAMPLE:
+//
+//     Given a union like `str | tuple[int] | tuple[int, int] | tuple[str, str]`, we would build the
+//     following instances of this struct:
+//
+//         HomogenousTupleGroup{
+//             element_type: int,
+//             original_tuple_types: [tuple[int], tuple[int, int]]
+//         }
+//
+//         HomogenousTupleGroup{
+//             element_type: str,
+//             original_tuple_types: [tuple[str, str]]
+//         }
+struct HomogeneousTupleGroup<'db> {
+    element_type: Type<'db>,
+    original_tuple_types: Vec<Type<'db>>,
+}
+
+impl<'db> HomogeneousTupleGroup<'db> {
+    fn new(element_type: Type<'db>, original_tuple_type: Type<'db>) -> Self {
+        Self {
+            element_type,
+            original_tuple_types: vec![original_tuple_type],
+        }
+    }
+
+    // Add a new tuple to this group.
+    fn add(&mut self, original_tuple_type: Type<'db>) {
+        self.original_tuple_types.push(original_tuple_type);
+    }
+
+    // Does the group have tuples of different lengths?
+    //
+    // This determines whether or not we should promote the types in this group to a single tuple
+    // of variable-length.
+    fn has_multiple_lengths(&self, db: &'db dyn Db) -> bool {
+        let mut lengths = self.original_tuple_types.iter().filter_map(|tuple_type| {
+            tuple_type
+                .homogeneous_fixed_length_tuple_instance(db)
+                .map(|(_, length)| length)
+        });
+
+        let Some(first_length) = lengths.next() else {
+            return false;
+        };
+
+        lengths.any(|length| length != first_length)
+    }
+}
+
+fn partition_homogeneous_fixed_length_tuple_union_elements<'db>(
+    db: &'db dyn Db,
+    elements: impl IntoIterator<Item = Type<'db>>,
+) -> (Vec<Type<'db>>, Vec<HomogeneousTupleGroup<'db>>) {
+    let mut other_union_elements = Vec::new();
+    let mut tuple_groups: Vec<HomogeneousTupleGroup<'db>> = Vec::new();
+
+    for element in elements {
+        if let Some((tuple_element_type, _)) = element.homogeneous_fixed_length_tuple_instance(db) {
+            if let Some(group) = tuple_groups
+                .iter_mut()
+                .find(|group| group.element_type.is_equivalent_to(db, tuple_element_type))
+            {
+                group.add(element);
+            } else {
+                tuple_groups.push(HomogeneousTupleGroup::new(tuple_element_type, element));
+            }
+        } else {
+            other_union_elements.push(element);
+        }
+    }
+
+    (other_union_elements, tuple_groups)
+}
 
 #[derive(Debug, Clone, Copy)]
 #[expect(clippy::struct_field_names)]
@@ -1297,6 +1375,32 @@ impl<'db> Type<'db> {
             .and_then(|instance| instance.own_tuple_spec(db))
     }
 
+    /// Detects whether or not this type is a homogeneously-typed tuple of fixed length
+    /// e.g. `tuple[str, str]` but NOT `tuple[str, int]` or `tuple[str, ...]`.
+    ///
+    /// If this type is indeed a homogeneously-typed, fixed-length tuple, then returns the
+    /// tuple's homogeneous element type and its length.
+    pub(crate) fn homogeneous_fixed_length_tuple_instance(
+        self,
+        db: &'db dyn Db,
+    ) -> Option<(Type<'db>, usize)> {
+        let tuple_spec = self.exact_tuple_instance_spec(db)?;
+        let TupleSpec::Fixed(tuple) = tuple_spec.as_ref() else {
+            return None;
+        };
+
+        let length = tuple.len();
+        if length == 0 {
+            return None;
+        }
+
+        let mut elements = tuple.iter_all_elements();
+        let element_type = elements.next()?;
+        elements
+            .all(|element| element.is_equivalent_to(db, element_type))
+            .then_some((element_type, length))
+    }
+
     /// Returns the materialization of this type depending on the given `variance`.
     ///
     /// More concretely, `T'`, the materialization of `T`, is the type `T` with all occurrences of
@@ -1925,6 +2029,78 @@ impl<'db> Type<'db> {
         )
     }
 
+    /// Promote unions of homogeneously-typed, fixed-length tuples with different lengths to a
+    /// single variable-length tuple when every tuple in that homogeneous group came from a tuple
+    /// literal in the collection literal currently being inferred.
+    ///
+    /// This deliberately only applies to unions; a standalone fixed-length tuple keeps its shape,
+    /// and groups that include non-literal tuple members remain unchanged.
+    ///
+    /// EXAMPLE:
+    ///
+    /// In the code below, we promote `dict[str, tuple[str, str] | tuple[str, str, str, str]]`
+    /// to `dict[str, tuple[str,...]]`.
+    ///
+    ///     languages = {
+    ///         "python": (".py", ".pyi"),
+    ///         "javascript": (".js", ".jsx", ".ts", ".tsx"),
+    ///     }
+    ///     reveal_type(languages)  # revealed: dict[str, tuple[str, ...]]
+    ///
+    pub(crate) fn promote_tuple_literal_unions(
+        self,
+        db: &'db dyn Db,
+        tuple_literal_candidates: &FxHashSet<Type<'db>>,
+        non_literal_tuple_candidates: &FxHashSet<Type<'db>>,
+    ) -> Type<'db> {
+        let Type::Union(union) = self else {
+            return self;
+        };
+
+        let (other_union_elements, tuple_groups) =
+            partition_homogeneous_fixed_length_tuple_union_elements(
+                db,
+                union.elements(db).iter().copied(),
+            );
+
+        if !tuple_groups.iter().any(|group| {
+            group.has_multiple_lengths(db)
+                && group.original_tuple_types.iter().all(|tuple_type| {
+                    tuple_literal_candidates.contains(tuple_type)
+                        && !non_literal_tuple_candidates.contains(tuple_type)
+                })
+        }) {
+            // No tuple group consisted entirely of tuple-literal members with differing lengths,
+            // so there is nothing to promote. Return early to avoid rebuilding the union type.
+            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(db)
+                && group.original_tuple_types.iter().all(|tuple_type| {
+                    tuple_literal_candidates.contains(tuple_type)
+                        && !non_literal_tuple_candidates.contains(tuple_type)
+                })
+            {
+                builder = builder.add(Type::homogeneous_tuple(db, group.element_type));
+            } else {
+                for element in group.original_tuple_types {
+                    builder = builder.add(element);
+                }
+            }
+        }
+
+        builder.build()
+    }
+
     /// Promote a top-level singleton type (like `None`, `EllipsisType`) to `T | Unknown`.
     pub(crate) fn promote_singletons(self, db: &'db dyn Db) -> Type<'db> {
         self.promote_singletons_impl(db)