return tuple spec when iterating

dcreager · dcreager · commit cfb020dcc784 · 2025-07-21T17:05:47.000-04:00
diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs
@@ -56,7 +56,7 @@ use crate::types::infer::infer_unpack_types;
 use crate::types::mro::{Mro, MroError, MroIterator};
 pub(crate) use crate::types::narrow::infer_narrowing_constraint;
 use crate::types::signatures::{Parameter, ParameterForm, Parameters, walk_signature};
-use crate::types::tuple::TupleType;
+use crate::types::tuple::{TupleSpec, TupleType};
 pub use crate::util::diagnostics::add_inferred_python_version_hint_to_diagnostic;
 use crate::{Db, FxOrderSet, Module, Program};
 pub(crate) use class::{ClassLiteral, ClassType, GenericAlias, KnownClass};
@@ -4610,35 +4610,43 @@ impl<'db> Type<'db> {
         }
     }
 
-    /// Returns the element type when iterating over `self`.
+    /// Returns a tuple spec describing the elements that are produced when iterating over `self`.
     ///
     /// This method should only be used outside of type checking because it omits any errors.
     /// For type checking, use [`try_iterate`](Self::try_iterate) instead.
-    fn iterate(self, db: &'db dyn Db) -> Type<'db> {
+    fn iterate(self, db: &'db dyn Db) -> TupleSpec<'db> {
         self.try_iterate(db)
-            .unwrap_or_else(|err| err.fallback_element_type(db))
+            .unwrap_or_else(|err| TupleSpec::homogeneous(err.fallback_element_type(db)))
     }
 
     /// Given the type of an object that is iterated over in some way,
-    /// return the type of objects that are yielded by that iteration.
+    /// return a tuple spec describing the type of objects that are yielded by that iteration.
     ///
-    /// E.g., for the following loop, given the type of `x`, infer the type of `y`:
+    /// E.g., for the following call, given the type of `x`, infer the types of the values that are
+    /// splatted into `y`'s positional arguments:
+    /// of `y`:
     /// ```python
-    /// for y in x:
-    ///     pass
+    /// y(*x)
     /// ```
-    fn try_iterate(self, db: &'db dyn Db) -> Result<Type<'db>, IterationError<'db>> {
-        if let Type::Tuple(tuple_type) = self {
-            return Ok(UnionType::from_elements(
-                db,
-                tuple_type.tuple(db).all_elements(),
-            ));
-        }
-
-        if let Type::GenericAlias(alias) = self {
-            if alias.origin(db).is_known(db, KnownClass::Tuple) {
-                return Ok(todo_type!("*tuple[] annotations"));
-            }
+    fn try_iterate(self, db: &'db dyn Db) -> Result<TupleSpec<'db>, IterationError<'db>> {
+        match self {
+            // XXX: Cow?
+            Type::Tuple(tuple_type) => return Ok(tuple_type.tuple(db).clone()),
+            Type::GenericAlias(alias) if alias.origin(db).is_known(db, KnownClass::Tuple) => {
+                return Ok(TupleSpec::homogeneous(todo_type!("*tuple[] annotations")));
+            }
+            Type::StringLiteral(string_literal_ty) => {
+                // We could go further and deconstruct to an array of `StringLiteral`
+                // with each individual character, instead of just an array of
+                // `LiteralString`, but there would be a cost and it's not clear that
+                // it's worth it.
+                return Ok(TupleSpec::from_elements(std::iter::repeat_n(
+                    Type::LiteralString,
+                    string_literal_ty.python_len(db),
+                )));
+            }
+            Type::LiteralString => return Ok(TupleSpec::homogeneous(Type::LiteralString)),
+            _ => {}
         }
 
         let try_call_dunder_getitem = || {
@@ -4664,12 +4672,14 @@ impl<'db> Type<'db> {
             Ok(iterator) => {
                 // `__iter__` is definitely bound and calling it succeeds.
                 // See what calling `__next__` on the object returned by `__iter__` gives us...
-                try_call_dunder_next_on_iterator(iterator).map_err(|dunder_next_error| {
-                    IterationError::IterReturnsInvalidIterator {
-                        iterator,
-                        dunder_next_error,
-                    }
-                })
+                try_call_dunder_next_on_iterator(iterator)
+                    .map(TupleSpec::homogeneous)
+                    .map_err(
+                        |dunder_next_error| IterationError::IterReturnsInvalidIterator {
+                            iterator,
+                            dunder_next_error,
+                        },
+                    )
             }
 
             // `__iter__` is possibly unbound...
@@ -4687,10 +4697,10 @@ impl<'db> Type<'db> {
                                 // and the type returned by the `__getitem__` method.
                                 //
                                 // No diagnostic is emitted; iteration will always succeed!
-                                UnionType::from_elements(
+                                TupleSpec::homogeneous(UnionType::from_elements(
                                     db,
                                     [dunder_next_return, dunder_getitem_return_type],
-                                )
+                                ))
                             })
                             .map_err(|dunder_getitem_error| {
                                 IterationError::PossiblyUnboundIterAndGetitemError {
@@ -4713,13 +4723,13 @@ impl<'db> Type<'db> {
             }
 
             // There's no `__iter__` method. Try `__getitem__` instead...
-            Err(CallDunderError::MethodNotAvailable) => {
-                try_call_dunder_getitem().map_err(|dunder_getitem_error| {
-                    IterationError::UnboundIterAndGetitemError {
+            Err(CallDunderError::MethodNotAvailable) => try_call_dunder_getitem()
+                .map(TupleSpec::homogeneous)
+                .map_err(
+                    |dunder_getitem_error| IterationError::UnboundIterAndGetitemError {
                         dunder_getitem_error,
-                    }
-                })
-            }
+                    },
+                ),
         }
     }
 
diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs
@@ -24,7 +24,7 @@ use crate::types::function::{
 };
 use crate::types::generics::{Specialization, SpecializationBuilder, SpecializationError};
 use crate::types::signatures::{Parameter, ParameterForm, Parameters};
-use crate::types::tuple::TupleType;
+use crate::types::tuple::{TupleSpec, TupleType};
 use crate::types::{
     BoundMethodType, ClassLiteral, DataclassParams, KnownClass, KnownInstanceType,
     MethodWrapperKind, PropertyInstanceType, SpecialFormType, TypeMapping, UnionType,
@@ -969,15 +969,15 @@ impl<'db> Bindings<'db> {
                                         // but `tuple[Never, ...]` eagerly simplifies to `tuple[()]`,
                                         // which will cause us to emit false positives if we index into the tuple.
                                         // Using `tuple[Unknown, ...]` avoids these false positives.
-                                        let specialization = if argument.is_never() {
-                                            Type::unknown()
+                                        let tuple_spec = if argument.is_never() {
+                                            TupleSpec::homogeneous(Type::unknown())
                                         } else {
                                             argument.try_iterate(db).expect(
                                                 "try_iterate() should not fail on a type \
                                                     assignable to `Iterable`",
                                             )
                                         };
-                                        TupleType::homogeneous(db, specialization)
+                                        Type::tuple(TupleType::new(db, tuple_spec))
                                     });
                                 overload.set_return_type(overridden_return);
                             }
diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs
@@ -2057,7 +2057,8 @@ impl<'db> ClassLiteral<'db> {
                                     index.expression(for_stmt.iterable(&module)),
                                 );
                                 // TODO: Potential diagnostics resulting from the iterable are currently not reported.
-                                let inferred_ty = iterable_ty.iterate(db);
+                                let inferred_ty =
+                                    iterable_ty.iterate(db).homogeneous_element_type(db);
 
                                 union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
                             }
@@ -2115,7 +2116,8 @@ impl<'db> ClassLiteral<'db> {
                                     index.expression(comprehension.iterable(&module)),
                                 );
                                 // TODO: Potential diagnostics resulting from the iterable are currently not reported.
-                                let inferred_ty = iterable_ty.iterate(db);
+                                let inferred_ty =
+                                    iterable_ty.iterate(db).homogeneous_element_type(db);
 
                                 union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
                             }
diff --git a/crates/ty_python_semantic/src/types/generics.rs b/crates/ty_python_semantic/src/types/generics.rs
@@ -194,7 +194,7 @@ impl<'db> GenericContext<'db> {
         db: &'db dyn Db,
         tuple: TupleType<'db>,
     ) -> Specialization<'db> {
-        let element_type = UnionType::from_elements(db, tuple.tuple(db).all_elements());
+        let element_type = tuple.tuple(db).homogeneous_element_type(db);
         Specialization::new(db, self, Box::from([element_type]), Some(tuple))
     }
 
diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs
@@ -4226,6 +4226,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             builder
                 .infer_standalone_expression(iter_expr)
                 .iterate(builder.db())
+                .homogeneous_element_type(builder.db())
         });
 
         self.infer_body(body);
@@ -4255,10 +4256,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 }
                 TargetKind::Single => {
                     let iterable_type = self.infer_standalone_expression(iterable);
-                    iterable_type.try_iterate(self.db()).unwrap_or_else(|err| {
-                        err.report_diagnostic(&self.context, iterable_type, iterable.into());
-                        err.fallback_element_type(self.db())
-                    })
+                    iterable_type
+                        .try_iterate(self.db())
+                        .map(|tuple| tuple.homogeneous_element_type(self.db()))
+                        .unwrap_or_else(|err| {
+                            err.report_diagnostic(&self.context, iterable_type, iterable.into());
+                            err.fallback_element_type(self.db())
+                        })
                 }
             }
         };
@@ -5370,6 +5374,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 builder.infer_standalone_expression(iter_expr)
             }
             .iterate(builder.db())
+            .homogeneous_element_type(builder.db())
         });
         for expr in ifs {
             self.infer_standalone_expression(expr);
@@ -5421,10 +5426,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 }
                 TargetKind::Single => {
                     let iterable_type = infer_iterable_type();
-                    iterable_type.try_iterate(self.db()).unwrap_or_else(|err| {
-                        err.report_diagnostic(&self.context, iterable_type, iterable.into());
-                        err.fallback_element_type(self.db())
-                    })
+                    iterable_type
+                        .try_iterate(self.db())
+                        .map(|tuple| tuple.homogeneous_element_type(self.db()))
+                        .unwrap_or_else(|err| {
+                            err.report_diagnostic(&self.context, iterable_type, iterable.into());
+                            err.fallback_element_type(self.db())
+                        })
                 }
             }
         };
@@ -5753,10 +5761,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         } = starred;
 
         let iterable_type = self.infer_expression(value);
-        iterable_type.try_iterate(self.db()).unwrap_or_else(|err| {
-            err.report_diagnostic(&self.context, iterable_type, value.as_ref().into());
-            err.fallback_element_type(self.db())
-        });
+        iterable_type
+            .try_iterate(self.db())
+            .map(|tuple| tuple.homogeneous_element_type(self.db()))
+            .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")
@@ -5780,10 +5791,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         } = yield_from;
 
         let iterable_type = self.infer_expression(value);
-        iterable_type.try_iterate(self.db()).unwrap_or_else(|err| {
-            err.report_diagnostic(&self.context, iterable_type, value.as_ref().into());
-            err.fallback_element_type(self.db())
-        });
+        iterable_type
+            .try_iterate(self.db())
+            .map(|tuple| tuple.homogeneous_element_type(self.db()))
+            .unwrap_or_else(|err| {
+                err.report_diagnostic(&self.context, iterable_type, value.as_ref().into());
+                err.fallback_element_type(self.db())
+            });
 
         // TODO get type from `ReturnType` of generator
         todo_type!("Generic `typing.Generator` type")
diff --git a/crates/ty_python_semantic/src/types/tuple.rs b/crates/ty_python_semantic/src/types/tuple.rs
@@ -1006,6 +1006,10 @@ impl<T> Tuple<T> {
 }
 
 impl<'db> Tuple<Type<'db>> {
+    pub(crate) fn homogeneous_element_type(&self, db: &'db dyn Db) -> Type<'db> {
+        UnionType::from_elements(db, self.all_elements())
+    }
+
     /// Concatenates another tuple to the end of this tuple, returning a new tuple.
     pub(crate) fn concat(&self, db: &'db dyn Db, other: &Self) -> Self {
         match self {
diff --git a/crates/ty_python_semantic/src/types/unpacker.rs b/crates/ty_python_semantic/src/types/unpacker.rs
@@ -1,5 +1,3 @@
-use std::borrow::Cow;
-
 use ruff_db::parsed::ParsedModuleRef;
 use rustc_hash::FxHashMap;
 
@@ -8,7 +6,7 @@ use ruff_python_ast::{self as ast, AnyNodeRef};
 use crate::Db;
 use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
 use crate::semantic_index::place::ScopeId;
-use crate::types::tuple::{ResizeTupleError, Tuple, TupleLength, TupleUnpacker};
+use crate::types::tuple::{ResizeTupleError, Tuple, TupleLength, TupleSpec, TupleUnpacker};
 use crate::types::{Type, TypeCheckDiagnostics, infer_expression_types};
 use crate::unpack::{UnpackKind, UnpackValue};
 
@@ -64,14 +62,17 @@ impl<'db, 'ast> Unpacker<'db, 'ast> {
                     value_type
                 }
             }
-            UnpackKind::Iterable => value_type.try_iterate(self.db()).unwrap_or_else(|err| {
-                err.report_diagnostic(
-                    &self.context,
-                    value_type,
-                    value.as_any_node_ref(self.db(), self.module()),
-                );
-                err.fallback_element_type(self.db())
-            }),
+            UnpackKind::Iterable => value_type
+                .try_iterate(self.db())
+                .map(|tuple| tuple.homogeneous_element_type(self.db()))
+                .unwrap_or_else(|err| {
+                    err.report_diagnostic(
+                        &self.context,
+                        value_type,
+                        value.as_any_node_ref(self.db(), self.module()),
+                    );
+                    err.fallback_element_type(self.db())
+                }),
             UnpackKind::ContextManager => value_type.try_enter(self.db()).unwrap_or_else(|err| {
                 err.report_diagnostic(
                     &self.context,
@@ -118,32 +119,12 @@ impl<'db, 'ast> Unpacker<'db, 'ast> {
                 };
 
                 for ty in unpack_types.iter().copied() {
-                    let tuple = match ty {
-                        Type::Tuple(tuple_ty) => Cow::Borrowed(tuple_ty.tuple(self.db())),
-                        Type::StringLiteral(string_literal_ty) => {
-                            // We could go further and deconstruct to an array of `StringLiteral`
-                            // with each individual character, instead of just an array of
-                            // `LiteralString`, but there would be a cost and it's not clear that
-                            // it's worth it.
-                            Cow::Owned(Tuple::from_elements(std::iter::repeat_n(
-                                Type::LiteralString,
-                                string_literal_ty.python_len(self.db()),
-                            )))
-                        }
-                        Type::LiteralString => Cow::Owned(Tuple::homogeneous(Type::LiteralString)),
-                        _ => {
-                            // TODO: Update our iterator protocol machinery to return a tuple
-                            // describing the returned values in more detail, when we can.
-                            Cow::Owned(Tuple::homogeneous(
-                                ty.try_iterate(self.db()).unwrap_or_else(|err| {
-                                    err.report_diagnostic(&self.context, ty, value_expr);
-                                    err.fallback_element_type(self.db())
-                                }),
-                            ))
-                        }
-                    };
+                    let tuple = ty.try_iterate(self.db()).unwrap_or_else(|err| {
+                        err.report_diagnostic(&self.context, ty, value_expr);
+                        TupleSpec::homogeneous(err.fallback_element_type(self.db()))
+                    });
 
-                    if let Err(err) = unpacker.unpack_tuple(tuple.as_ref()) {
+                    if let Err(err) = unpacker.unpack_tuple(&tuple) {
                         unpacker
                             .unpack_tuple(&Tuple::homogeneous(Type::unknown()))
                             .expect("adding a homogeneous tuple should always succeed");

Original file line number	Diff line number	Diff line change
`@@ -194,7 +194,7 @@ impl<'db> GenericContext<'db> {`
`194`	`194`	`db: &'db dyn Db,`
`195`	`195`	`tuple: TupleType<'db>,`
`196`	`196`	`) -> Specialization<'db> {`
`197`		`- let element_type = UnionType::from_elements(db, tuple.tuple(db).all_elements());`
	`197`	`+ let element_type = tuple.tuple(db).homogeneous_element_type(db);`
`198`	`198`	`Specialization::new(db, self, Box::from([element_type]), Some(tuple))`
`199`	`199`	`}`
`200`	`200`