[ty] Fix Todo type for starred elements in tuple expressions

Author: AlexWaygood

crates/ty_python_semantic/resources/mdtest/bidirectional.md

@@ -42,6 +42,12 @@ def f[T](x: T, cond: bool) -> T | list[T]:
     return x if cond else [x]
 
 l5: int | list[int] = f(1, True)
+
+a: list[int] = [1, 2, *(3, 4, 5)]
+reveal_type(a)  # revealed: list[int]
+
+b: list[list[int]] = [[1], [2], *([3], [4])]
+reveal_type(b)  # revealed: list[list[int]]
 ```
 
 `typed_dict.py`:

crates/ty_python_semantic/resources/mdtest/expression/len.md

@@ -43,13 +43,13 @@ reveal_type(len((1,)))  # revealed: Literal[1]
 reveal_type(len((1, 2)))  # revealed: Literal[2]
 reveal_type(len(tuple()))  # revealed: Literal[0]
 
-# TODO: Handle star unpacks; Should be: Literal[0]
-reveal_type(len((*[],)))  # revealed: Literal[1]
+# could also be `Literal[0]`, but `int` is accurate
+reveal_type(len((*[],)))  # revealed: int
 
 # fmt: off
 
-# TODO: Handle star unpacks; Should be: Literal[1]
-reveal_type(len(  # revealed: Literal[2]
+# could also be `Literal[1]`, but `int` is accurate
+reveal_type(len(  # revealed: int
     (
         *[],
         1,
@@ -58,11 +58,11 @@ reveal_type(len(  # revealed: Literal[2]
 
 # fmt: on
 
-# TODO: Handle star unpacks; Should be: Literal[2]
-reveal_type(len((*[], 1, 2)))  # revealed: Literal[3]
+# Could also be `Literal[2]`, but `int` is accurate
+reveal_type(len((*[], 1, 2)))  # revealed: int
 
-# TODO: Handle star unpacks; Should be: Literal[0]
-reveal_type(len((*[], *{})))  # revealed: Literal[2]
+# Could also be `Literal[0]`, but `int` is accurate
+reveal_type(len((*[], *{})))  # revealed: int
 ```
 
 Tuple subclasses:

crates/ty_python_semantic/resources/mdtest/type_compendium/tuple.md

@@ -531,4 +531,18 @@ x: list[Literal[1, 2, 3]] = list((1, 2, 3))
 reveal_type(x)  # revealed: list[Literal[1, 2, 3]]
 ```
 
+## Tuples with starred elements
+
+```py
+x = (1, *range(3), 3)
+reveal_type(x)  # revealed: tuple[Literal[1], *tuple[int, ...], Literal[3]]
+
+y = 1, 2
+
+reveal_type(("foo", *y))  # revealed: tuple[Literal["foo"], Literal[1], Literal[2]]
+
+aa: tuple[list[int], ...] = ([42], *{[56], [78]}, [100])
+reveal_type(aa)  # revealed: tuple[list[int], *tuple[list[int], ...], list[int]]
+```
+
 [not a singleton type]: https://discuss.python.org/t/should-we-specify-in-the-language-reference-that-the-empty-tuple-is-a-singleton/67957

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

@@ -2941,7 +2941,11 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
     ) {
         let parameters = self.signature.parameters();
         let parameter = &parameters[parameter_index];
-        if let Some(mut expected_ty) = parameter.annotated_type() {
+
+        // TODO: handle starred annotations, e.g. `*args: *Ts` or `*args: *tuple[int, *tuple[str, ...]]`
+        if let Some(mut expected_ty) = parameter.annotated_type()
+            && !parameter.has_starred_annotation()
+        {
             if let Some(specialization) = self.specialization {
                 argument_type = argument_type.apply_specialization(self.db, specialization);
                 expected_ty = expected_ty.apply_specialization(self.db, specialization);

crates/ty_python_semantic/src/types/infer.rs

@@ -377,6 +377,27 @@ impl<'db> TypeContext<'db> {
             annotation: self.annotation.map(f),
         }
     }
+
+    pub(crate) fn for_starred_expression(
+        db: &'db dyn Db,
+        expected_element_type: Type<'db>,
+        expr: &ast::ExprStarred,
+    ) -> Self {
+        match &*expr.value {
+            ast::Expr::List(_) => Self::new(Some(
+                KnownClass::List.to_specialized_instance(db, [expected_element_type]),
+            )),
+            ast::Expr::Set(_) => Self::new(Some(
+                KnownClass::Set.to_specialized_instance(db, [expected_element_type]),
+            )),
+            ast::Expr::Tuple(_) => {
+                Self::new(Some(Type::homogeneous_tuple(db, expected_element_type)))
+            }
+            // `Iterable[<expected_element_type>]` would work well for an arbitrary other node
+            // if <https://github.com/astral-sh/ty/issues/1576> is implemented.
+            _ => Self::default(),
+        }
+    }
 }
 
 /// Returns the statically-known truthiness of a given expression.

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -94,7 +94,9 @@ use crate::types::mro::MroErrorKind;
 use crate::types::newtype::NewType;
 use crate::types::signatures::Signature;
 use crate::types::subclass_of::SubclassOfInner;
-use crate::types::tuple::{Tuple, TupleLength, TupleSpec, TupleType};
+use crate::types::tuple::{
+    Tuple, TupleLength, TupleSpec, TupleSpecBuilder, TupleType, VariableLengthTuple,
+};
 use crate::types::typed_dict::{
     TypedDictAssignmentKind, validate_typed_dict_constructor, validate_typed_dict_dict_literal,
     validate_typed_dict_key_assignment,
@@ -6926,7 +6928,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             ast::Expr::If(if_expression) => self.infer_if_expression(if_expression, tcx),
             ast::Expr::Lambda(lambda_expression) => self.infer_lambda_expression(lambda_expression),
             ast::Expr::Call(call_expression) => self.infer_call_expression(call_expression, tcx),
-            ast::Expr::Starred(starred) => self.infer_starred_expression(starred),
+            ast::Expr::Starred(starred) => self.infer_starred_expression(starred, tcx),
             ast::Expr::Yield(yield_expression) => self.infer_yield_expression(yield_expression),
             ast::Expr::YieldFrom(yield_from) => self.infer_yield_from_expression(yield_from),
             ast::Expr::Await(await_expression) => self.infer_await_expression(await_expression),
@@ -7151,25 +7153,66 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             )
         });
 
+        let mut is_homogeneous_tuple_annotation = false;
+
         let annotated_tuple = tcx
             .known_specialization(self.db(), KnownClass::Tuple)
             .and_then(|specialization| {
-                specialization
+                let spec = specialization
                     .tuple(self.db())
-                    .expect("the specialization of `KnownClass::Tuple` must have a tuple spec")
-                    .resize(self.db(), TupleLength::Fixed(elts.len()))
-                    .ok()
+                    .expect("the specialization of `KnownClass::Tuple` must have a tuple spec");
+
+                if matches!(
+                    spec,
+                    Tuple::Variable(VariableLengthTuple { prefix, variable: _, suffix})
+                    if prefix.is_empty() && suffix.is_empty()
+                ) {
+                    is_homogeneous_tuple_annotation = true;
+                }
+
+                spec.resize(self.db(), TupleLength::Fixed(elts.len())).ok()
             });
 
         let mut annotated_elt_tys = annotated_tuple.as_ref().map(Tuple::all_elements);
 
         let db = self.db();
-        let element_types = elts.iter().map(|element| {
-            let annotated_elt_ty = annotated_elt_tys.as_mut().and_then(Iterator::next).copied();
-            self.infer_expression(element, TypeContext::new(annotated_elt_ty))
-        });
 
-        Type::heterogeneous_tuple(db, element_types)
+        let can_use_type_context =
+            is_homogeneous_tuple_annotation || elts.iter().all(|elt| !elt.is_starred_expr());
+
+        let mut infer_element = |elt: &ast::Expr| {
+            if can_use_type_context {
+                let annotated_elt_ty = annotated_elt_tys.as_mut().and_then(Iterator::next).copied();
+                let context = if let ast::Expr::Starred(starred) = elt {
+                    annotated_elt_ty
+                        .map(|expected_element_type| {
+                            TypeContext::for_starred_expression(db, expected_element_type, starred)
+                        })
+                        .unwrap_or_default()
+                } else {
+                    TypeContext::new(annotated_elt_ty)
+                };
+                self.infer_expression(elt, context)
+            } else {
+                self.infer_expression(elt, TypeContext::default())
+            }
+        };
+
+        let mut builder = TupleSpecBuilder::with_capacity(elts.len());
+
+        for element in elts {
+            if element.is_starred_expr() {
+                let element_type = infer_element(element);
+                // Fine to use `iterate` rather than `try_iterate` here:
+                // errors from iterating over something not iterable will have been
+                // emitted in the `infer_element` call above.
+                builder = builder.concat(db, &element_type.iterate(db));
+            } else {
+                builder.push(infer_element(element).fallback_to_divergent(db));
+            }
+        }
+
+        Type::tuple(TupleType::new(db, &builder.build()))
     }
 
     fn infer_list_expression(&mut self, list: &ast::ExprList, tcx: TypeContext<'db>) -> Type<'db> {
@@ -7326,7 +7369,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         let inferable = generic_context.inferable_typevars(self.db());
 
-        // Remove any union elements of that are unrelated to the collection type.
+        // Remove any union elements of the annotation that are unrelated to the collection type.
         //
         // For example, we only want the `list[int]` from `annotation: list[int] | None` if
         // `collection_ty` is `list`.
@@ -7366,8 +7409,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         }
 
         let elt_tcxs = match annotated_elt_tys {
-            None => Either::Left(iter::repeat(TypeContext::default())),
-            Some(tys) => Either::Right(tys.iter().map(|ty| TypeContext::new(Some(*ty)))),
+            None => Either::Left(iter::repeat(None)),
+            Some(tys) => Either::Right(tys.iter().copied().map(Some)),
         };
 
         for elts in elts {
@@ -7396,6 +7439,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             {
                 let Some(elt) = elt else { continue };
 
+                let elt_tcx = if let ast::Expr::Starred(starred) = elt {
+                    elt_tcx
+                        .map(|ty| TypeContext::for_starred_expression(self.db(), ty, starred))
+                        .unwrap_or_default()
+                } else {
+                    TypeContext::new(elt_tcx)
+                };
+
                 let inferred_elt_ty = infer_elt_expression(self, elt, elt_tcx);
 
                 // Simplify the inference based on the declared type of the element.
@@ -7409,7 +7460,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 // unions for large nested list literals, which the constraint solver struggles with.
                 let inferred_elt_ty = inferred_elt_ty.promote_literals(self.db(), elt_tcx);
 
-                builder.infer(Type::TypeVar(elt_ty), inferred_elt_ty).ok()?;
+                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
+                        },
+                    )
+                    .ok()?;
             }
         }
 
@@ -8204,25 +8266,28 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         }
     }
 
-    fn infer_starred_expression(&mut self, starred: &ast::ExprStarred) -> Type<'db> {
+    fn infer_starred_expression(
+        &mut self,
+        starred: &ast::ExprStarred,
+        tcx: TypeContext<'db>,
+    ) -> Type<'db> {
         let ast::ExprStarred {
             range: _,
             node_index: _,
             value,
             ctx: _,
         } = starred;
 
-        let iterable_type = self.infer_expression(value, TypeContext::default());
+        let db = self.db();
+        let iterable_type = self.infer_expression(value, tcx);
+
         iterable_type
-            .try_iterate(self.db())
-            .map(|tuple| tuple.homogeneous_element_type(self.db()))
+            .try_iterate(db)
+            .map(|spec| Type::tuple(TupleType::new(db, &spec)))
             .unwrap_or_else(|err| {
                 err.report_diagnostic(&self.context, iterable_type, value.as_ref().into());
-                err.fallback_element_type(self.db())
-            });
-
-        // TODO
-        todo_type!("starred expression")
+                Type::homogeneous_tuple(db, err.fallback_element_type(db))
+            })
     }
 
     fn infer_yield_expression(&mut self, yield_expression: &ast::ExprYield) -> Type<'db> {

crates/ty_python_semantic/src/types/infer/builder/annotation_expression.rs

@@ -121,9 +121,9 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
             ast::Expr::StringLiteral(string) => self.infer_string_annotation_expression(string),
 
             // Annotation expressions also get special handling for `*args` and `**kwargs`.
-            ast::Expr::Starred(starred) => {
-                TypeAndQualifiers::declared(self.infer_starred_expression(starred))
-            }
+            ast::Expr::Starred(starred) => TypeAndQualifiers::declared(
+                self.infer_starred_expression(starred, TypeContext::default()),
+            ),
 
             ast::Expr::BytesLiteral(bytes) => {
                 if let Some(builder) = self