[ty] Support variable-length tuples in unpacking assignments (astral-sh#18948)

This PR updates our unpacking assignment logic to use the new tuple
machinery. As a result, we can now unpack variable-length tuples
correctly.

As part of this, the `TupleSpec` classes have been renamed to `Tuple`,
and can now contain any element (Rust) type, not just `Type<'db>`. The
unpacker uses a tuple of `UnionBuilder`s to maintain the types that will
be assigned to each target, as we iterate through potentially many union
elements on the rhs. We also add a new consuming iterator for tuples,
and update the `all_elements` methods to wrap the result in an enum
(similar to `itertools::Position`) letting you know which part of the
tuple each element appears in. I also added a new
`UnionBuilder::try_build`, which lets you specify a different fallback
type if the union contains no elements.

Author: dcreager

crates/ty_python_semantic/resources/mdtest/snapshots/unpacking.md_-_Unpacking_-_Too_few_values_to_un…_(cef19e6b2b58e6a3).snap

@@ -24,7 +24,7 @@ error[invalid-assignment]: Not enough values to unpack
 1 | [a, *b, c, d] = (1, 2)  # error: [invalid-assignment]
   | ^^^^^^^^^^^^^   ------ Got 2
   | |
-  | Expected 3 or more
+  | Expected at least 3
   |
 info: rule `invalid-assignment` is enabled by default
 

crates/ty_python_semantic/resources/mdtest/unpacking.md

@@ -106,7 +106,7 @@ reveal_type(d)  # revealed: Literal[5]
 ### Starred expression (1)
 
 ```py
-# error: [invalid-assignment] "Not enough values to unpack: Expected 3 or more"
+# error: [invalid-assignment] "Not enough values to unpack: Expected at least 3"
 [a, *b, c, d] = (1, 2)
 reveal_type(a)  # revealed: Unknown
 reveal_type(b)  # revealed: list[Unknown]
@@ -119,7 +119,7 @@ reveal_type(d)  # revealed: Unknown
 ```py
 [a, *b, c] = (1, 2)
 reveal_type(a)  # revealed: Literal[1]
-reveal_type(b)  # revealed: list[Unknown]
+reveal_type(b)  # revealed: list[Never]
 reveal_type(c)  # revealed: Literal[2]
 ```
 
@@ -154,7 +154,7 @@ reveal_type(c)  # revealed: list[Literal[3, 4]]
 ### Starred expression (6)
 
 ```py
-# error: [invalid-assignment] "Not enough values to unpack: Expected 5 or more"
+# error: [invalid-assignment] "Not enough values to unpack: Expected at least 5"
 (a, b, c, *d, e, f) = (1,)
 reveal_type(a)  # revealed: Unknown
 reveal_type(b)  # revealed: Unknown
@@ -258,6 +258,155 @@ def _(value: list[int]):
     reveal_type(c)  # revealed: int
 ```
 
+## Homogeneous tuples
+
+### Simple unpacking
+
+```py
+def _(value: tuple[int, ...]):
+    a, b = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: int
+```
+
+### Nested unpacking
+
+```py
+def _(value: tuple[tuple[int, ...], ...]):
+    a, (b, c) = value
+    reveal_type(a)  # revealed: tuple[int, ...]
+    reveal_type(b)  # revealed: int
+    reveal_type(c)  # revealed: int
+```
+
+### Invalid nested unpacking
+
+```py
+def _(value: tuple[int, ...]):
+    # error: [not-iterable] "Object of type `int` is not iterable"
+    a, (b, c) = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: Unknown
+    reveal_type(c)  # revealed: Unknown
+```
+
+### Starred expression
+
+```py
+def _(value: tuple[int, ...]):
+    a, *b, c = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: list[int]
+    reveal_type(c)  # revealed: int
+```
+
+## Mixed tuples
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+### Simple unpacking (1)
+
+```py
+def _(value: tuple[int, *tuple[str, ...]]):
+    a, b = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: str
+```
+
+### Simple unpacking (2)
+
+```py
+def _(value: tuple[int, int, *tuple[str, ...]]):
+    a, b = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: int
+```
+
+### Simple unpacking (3)
+
+```py
+def _(value: tuple[int, *tuple[str, ...], int]):
+    a, b, c = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: str
+    reveal_type(c)  # revealed: int
+```
+
+### Invalid unpacked
+
+```py
+def _(value: tuple[int, int, int, *tuple[str, ...]]):
+    # error: [invalid-assignment] "Too many values to unpack: Expected 2"
+    a, b = value
+    reveal_type(a)  # revealed: Unknown
+    reveal_type(b)  # revealed: Unknown
+```
+
+### Nested unpacking
+
+```py
+def _(value: tuple[str, *tuple[tuple[int, ...], ...]]):
+    a, (b, c) = value
+    reveal_type(a)  # revealed: str
+    reveal_type(b)  # revealed: int
+    reveal_type(c)  # revealed: int
+```
+
+### Invalid nested unpacking
+
+```py
+def _(value: tuple[str, *tuple[int, ...]]):
+    # error: [not-iterable] "Object of type `int` is not iterable"
+    a, (b, c) = value
+    reveal_type(a)  # revealed: str
+    reveal_type(b)  # revealed: Unknown
+    reveal_type(c)  # revealed: Unknown
+```
+
+### Starred expression (1)
+
+```py
+def _(value: tuple[int, *tuple[str, ...]]):
+    a, *b, c = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: list[str]
+    reveal_type(c)  # revealed: str
+```
+
+### Starred expression (2)
+
+```py
+def _(value: tuple[int, *tuple[str, ...], int]):
+    a, *b, c = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: list[str]
+    reveal_type(c)  # revealed: int
+```
+
+### Starred expression (3)
+
+```py
+def _(value: tuple[int, *tuple[str, ...], int]):
+    a, *b, c, d = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: list[str]
+    reveal_type(c)  # revealed: str
+    reveal_type(d)  # revealed: int
+```
+
+### Starred expression (4)
+
+```py
+def _(value: tuple[int, int, *tuple[str, ...], int]):
+    a, *b, c = value
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: list[int | str]
+    reveal_type(c)  # revealed: int
+```
+
 ## String
 
 ### Simple unpacking
@@ -290,7 +439,7 @@ reveal_type(b)  # revealed: Unknown
 ### Starred expression (1)
 
 ```py
-# error: [invalid-assignment] "Not enough values to unpack: Expected 3 or more"
+# error: [invalid-assignment] "Not enough values to unpack: Expected at least 3"
 (a, *b, c, d) = "ab"
 reveal_type(a)  # revealed: Unknown
 reveal_type(b)  # revealed: list[Unknown]
@@ -299,7 +448,7 @@ reveal_type(d)  # revealed: Unknown
 ```
 
 ```py
-# error: [invalid-assignment] "Not enough values to unpack: Expected 3 or more"
+# error: [invalid-assignment] "Not enough values to unpack: Expected at least 3"
 (a, b, *c, d) = "a"
 reveal_type(a)  # revealed: Unknown
 reveal_type(b)  # revealed: Unknown
@@ -312,7 +461,7 @@ reveal_type(d)  # revealed: Unknown
 ```py
 (a, *b, c) = "ab"
 reveal_type(a)  # revealed: LiteralString
-reveal_type(b)  # revealed: list[Unknown]
+reveal_type(b)  # revealed: list[Never]
 reveal_type(c)  # revealed: LiteralString
 ```
 

crates/ty_python_semantic/src/types.rs

@@ -726,7 +726,7 @@ impl<'db> Type<'db> {
                         .map(|ty| ty.materialize(db, variance.flip())),
                 )
                 .build(),
-            Type::Tuple(tuple_type) => Type::tuple(db, tuple_type.materialize(db, variance)),
+            Type::Tuple(tuple_type) => Type::tuple(tuple_type.materialize(db, variance)),
             Type::TypeVar(type_var) => Type::TypeVar(type_var.materialize(db, variance)),
             Type::TypeIs(type_is) => {
                 type_is.with_type(db, type_is.return_type(db).materialize(db, variance))
@@ -1141,7 +1141,7 @@ impl<'db> Type<'db> {
         match self {
             Type::Union(union) => Type::Union(union.normalized(db)),
             Type::Intersection(intersection) => Type::Intersection(intersection.normalized(db)),
-            Type::Tuple(tuple) => Type::tuple(db, tuple.normalized(db)),
+            Type::Tuple(tuple) => Type::tuple(tuple.normalized(db)),
             Type::Callable(callable) => Type::Callable(callable.normalized(db)),
             Type::ProtocolInstance(protocol) => protocol.normalized(db),
             Type::NominalInstance(instance) => Type::NominalInstance(instance.normalized(db)),
@@ -3458,7 +3458,7 @@ impl<'db> Type<'db> {
             Type::BooleanLiteral(bool) => Truthiness::from(*bool),
             Type::StringLiteral(str) => Truthiness::from(!str.value(db).is_empty()),
             Type::BytesLiteral(bytes) => Truthiness::from(!bytes.value(db).is_empty()),
-            Type::Tuple(tuple) => match tuple.tuple(db).size_hint() {
+            Type::Tuple(tuple) => match tuple.tuple(db).len().size_hint() {
                 // The tuple type is AlwaysFalse if it contains only the empty tuple
                 (_, Some(0)) => Truthiness::AlwaysFalse,
                 // The tuple type is AlwaysTrue if its inhabitants must always have length >=1
@@ -4312,7 +4312,7 @@ impl<'db> Type<'db> {
                     let mut parameter =
                         Parameter::positional_only(Some(Name::new_static("iterable")))
                             .with_annotated_type(instantiated);
-                    if matches!(spec.size_hint().1, Some(0)) {
+                    if matches!(spec.len().maximum(), Some(0)) {
                         parameter = parameter.with_default_type(TupleType::empty(db));
                     }
                     Parameters::new([parameter])
@@ -5350,7 +5350,7 @@ impl<'db> Type<'db> {
                 }
                 builder.build()
             }
-            Type::Tuple(tuple) => Type::Tuple(tuple.apply_type_mapping(db, type_mapping)),
+            Type::Tuple(tuple) => Type::tuple(tuple.apply_type_mapping(db, type_mapping)),
 
             Type::TypeIs(type_is) => type_is.with_type(db, type_is.return_type(db).apply_type_mapping(db, type_mapping)),
 

crates/ty_python_semantic/src/types/builder.rs

@@ -444,6 +444,10 @@ impl<'db> UnionBuilder<'db> {
     }
 
     pub(crate) fn build(self) -> Type<'db> {
+        self.try_build().unwrap_or(Type::Never)
+    }
+
+    pub(crate) fn try_build(self) -> Option<Type<'db>> {
         let mut types = vec![];
         for element in self.elements {
             match element {
@@ -460,9 +464,12 @@ impl<'db> UnionBuilder<'db> {
             }
         }
         match types.len() {
-            0 => Type::Never,
-            1 => types[0],
-            _ => Type::Union(UnionType::new(self.db, types.into_boxed_slice())),
+            0 => None,
+            1 => Some(types[0]),
+            _ => Some(Type::Union(UnionType::new(
+                self.db,
+                types.into_boxed_slice(),
+            ))),
         }
     }
 }

crates/ty_python_semantic/src/types/call/arguments.rs

@@ -221,10 +221,10 @@ fn expand_type<'db>(db: &'db dyn Db, ty: Type<'db>) -> Option<Vec<Type<'db>>> {
             let expanded = tuple
                 .all_elements()
                 .map(|element| {
-                    if let Some(expanded) = expand_type(db, element) {
+                    if let Some(expanded) = expand_type(db, *element) {
                         Either::Left(expanded.into_iter())
                     } else {
-                        Either::Right(std::iter::once(element))
+                        Either::Right(std::iter::once(*element))
                     }
                 })
                 .multi_cartesian_product()

crates/ty_python_semantic/src/types/generics.rs

@@ -286,9 +286,13 @@ impl<'db> Specialization<'db> {
             return tuple;
         }
         if let [element_type] = self.types(db) {
-            return TupleType::new(db, TupleSpec::homogeneous(*element_type)).tuple(db);
+            if let Some(tuple) = TupleType::new(db, TupleSpec::homogeneous(*element_type)) {
+                return tuple.tuple(db);
+            }
         }
-        TupleType::new(db, TupleSpec::homogeneous(Type::unknown())).tuple(db)
+        TupleType::new(db, TupleSpec::homogeneous(Type::unknown()))
+            .expect("tuple[Unknown, ...] should never contain Never")
+            .tuple(db)
     }
 
     /// Returns the type that a typevar is mapped to, or None if the typevar isn't part of this
@@ -330,7 +334,7 @@ impl<'db> Specialization<'db> {
             .collect();
         let tuple_inner = self
             .tuple_inner(db)
-            .map(|tuple| tuple.apply_type_mapping(db, type_mapping));
+            .and_then(|tuple| tuple.apply_type_mapping(db, type_mapping));
         Specialization::new(db, self.generic_context(db), types, tuple_inner)
     }
 
@@ -374,7 +378,7 @@ impl<'db> Specialization<'db> {
 
     pub(crate) fn normalized(self, db: &'db dyn Db) -> Self {
         let types: Box<[_]> = self.types(db).iter().map(|ty| ty.normalized(db)).collect();
-        let tuple_inner = self.tuple_inner(db).map(|tuple| tuple.normalized(db));
+        let tuple_inner = self.tuple_inner(db).and_then(|tuple| tuple.normalized(db));
         Self::new(db, self.generic_context(db), types, tuple_inner)
     }
 
@@ -394,7 +398,7 @@ impl<'db> Specialization<'db> {
                 vartype.materialize(db, variance)
             })
             .collect();
-        let tuple_inner = self.tuple_inner(db).map(|tuple| {
+        let tuple_inner = self.tuple_inner(db).and_then(|tuple| {
             // Tuples are immutable, so tuple element types are always in covariant position.
             tuple.materialize(db, variance)
         });
@@ -637,7 +641,7 @@ impl<'db> SpecializationBuilder<'db> {
                     (TupleSpec::Fixed(formal_tuple), TupleSpec::Fixed(actual_tuple)) => {
                         if formal_tuple.len() == actual_tuple.len() {
                             for (formal_element, actual_element) in formal_tuple.elements().zip(actual_tuple.elements()) {
-                                self.infer(formal_element, actual_element)?;
+                                self.infer(*formal_element, *actual_element)?;
                             }
                         }
                     }