Fix unpack_sequence to match CPython behavior

youknowone · youknowone · commit 99d74e9a55e1 · 2026-02-17T16:52:34.000+09:00
Add fast path for exact tuple/list (direct length check without
iterator). General path now iterates only size+1 elements,
fixing hang on infinite sequences like __getitem__ without
raising IndexError.

Error messages now include "got N" for tuple, list, and dict.
Remove 7 expectedFailure/skip markers from test_unpack.
diff --git a/Lib/test/test_unpack.py b/Lib/test/test_unpack.py
@@ -70,14 +70,14 @@
 
 Unpacking tuple of wrong size
 
-    >>> a, b = t # TODO: RUSTPYTHON # doctest: +EXPECTED_FAILURE
+    >>> a, b = t
     Traceback (most recent call last):
       ...
     ValueError: too many values to unpack (expected 2, got 3)
 
 Unpacking tuple of wrong size
 
-    >>> a, b = l # TODO: RUSTPYTHON # doctest: +EXPECTED_FAILURE
+    >>> a, b = l
     Traceback (most recent call last):
       ...
     ValueError: too many values to unpack (expected 2, got 3)
@@ -144,7 +144,7 @@
 
 Unpacking to an empty iterable should raise ValueError
 
-    >>> () = [42] # TODO: RUSTPYTHON # doctest: +EXPECTED_FAILURE
+    >>> () = [42]
     Traceback (most recent call last):
       ...
     ValueError: too many values to unpack (expected 0, got 1)
@@ -157,13 +157,13 @@
     Traceback (most recent call last):
       ...
     ValueError: too many values to unpack (expected 3)
-    >>> next(it) # TODO: RUSTPYTHON; Raise `StopIteration`  # doctest: +SKIP
+    >>> next(it)
     4
 
 Unpacking unbalanced dict
 
     >>> d = {4: 'four', 5: 'five', 6: 'six', 7: 'seven'}
-    >>> a, b, c = d # TODO: RUSTPYTHON; # doctest: +EXPECTED_FAILURE
+    >>> a, b, c = d
     Traceback (most recent call last):
       ...
     ValueError: too many values to unpack (expected 3, got 4)
@@ -176,7 +176,7 @@
     ...     def __getitem__(self, i):
     ...         return i*2
     ...
-    >>> x, y, z = LengthTooLong() # TODO: RUSTPYTHON; Hangs # doctest: +SKIP
+    >>> x, y, z = LengthTooLong()
     Traceback (most recent call last):
       ...
     ValueError: too many values to unpack (expected 3)
@@ -189,7 +189,7 @@
     ...     def __getitem__(self, i):
     ...         return i*2
     ...
-    >>> x, y, z = BadLength() # TODO: RUSTPYTHON; Hangs # doctest: +SKIP
+    >>> x, y, z = BadLength()
     Traceback (most recent call last):
       ...
     ValueError: too many values to unpack (expected 3)
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -3010,13 +3010,54 @@ impl ExecutingFrame<'_> {
         Ok(None)
     }
 
+    /// _PyEval_UnpackIterableStackRef
     fn unpack_sequence(&mut self, size: u32, vm: &VirtualMachine) -> FrameResult {
         let value = self.pop_value();
+        let size = size as usize;
+
+        // Fast path for exact tuple/list types (not subclasses) — check
+        // length directly without creating an iterator, matching
+        // UNPACK_SEQUENCE_TUPLE / UNPACK_SEQUENCE_LIST specializations.
+        let cls = value.class();
+        let fast_elements: Option<Vec<PyObjectRef>> = if cls.is(vm.ctx.types.tuple_type) {
+            Some(value.downcast_ref::<PyTuple>().unwrap().as_slice().to_vec())
+        } else if cls.is(vm.ctx.types.list_type) {
+            Some(
+                value
+                    .downcast_ref::<PyList>()
+                    .unwrap()
+                    .borrow_vec()
+                    .to_vec(),
+            )
+        } else {
+            None
+        };
+        if let Some(elements) = fast_elements {
+            return match elements.len().cmp(&size) {
+                core::cmp::Ordering::Equal => {
+                    self.state
+                        .stack
+                        .extend(elements.into_iter().rev().map(Some));
+                    Ok(None)
+                }
+                core::cmp::Ordering::Greater => Err(vm.new_value_error(format!(
+                    "too many values to unpack (expected {size}, got {})",
+                    elements.len()
+                ))),
+                core::cmp::Ordering::Less => Err(vm.new_value_error(format!(
+                    "not enough values to unpack (expected {size}, got {})",
+                    elements.len()
+                ))),
+            };
+        }
+
+        // General path — iterate up to `size + 1` elements to avoid
+        // consuming the entire iterator (fixes hang on infinite sequences).
         let not_iterable = value.class().slots.iter.load().is_none()
             && value
                 .get_class_attr(vm.ctx.intern_str("__getitem__"))
                 .is_none();
-        let elements: Vec<_> = value.try_to_value(vm).map_err(|e| {
+        let iter = PyIter::try_from_object(vm, value.clone()).map_err(|e| {
             if not_iterable && e.class().is(vm.ctx.exceptions.type_error) {
                 vm.new_type_error(format!(
                     "cannot unpack non-iterable {} object",
@@ -3026,24 +3067,48 @@ impl ExecutingFrame<'_> {
                 e
             }
         })?;
-        let msg = match elements.len().cmp(&(size as usize)) {
-            core::cmp::Ordering::Equal => {
-                // Wrap each element in Some() for Option<PyObjectRef> stack
+
+        let mut elements = Vec::with_capacity(size);
+        for _ in 0..size {
+            match iter.next(vm)? {
+                PyIterReturn::Return(item) => elements.push(item),
+                PyIterReturn::StopIteration(_) => {
+                    return Err(vm.new_value_error(format!(
+                        "not enough values to unpack (expected {size}, got {})",
+                        elements.len()
+                    )));
+                }
+            }
+        }
+
+        // Check that the iterator is exhausted.
+        match iter.next(vm)? {
+            PyIterReturn::Return(_) => {
+                // For exact dict types, show "got N" using the container's
+                // size (PyDict_Size). Exact tuple/list are handled by the
+                // fast path above and never reach here.
+                let msg = if value.class().is(vm.ctx.types.dict_type) {
+                    if let Ok(got) = value.length(vm) {
+                        if got > size {
+                            format!("too many values to unpack (expected {size}, got {got})")
+                        } else {
+                            format!("too many values to unpack (expected {size})")
+                        }
+                    } else {
+                        format!("too many values to unpack (expected {size})")
+                    }
+                } else {
+                    format!("too many values to unpack (expected {size})")
+                };
+                Err(vm.new_value_error(msg))
+            }
+            PyIterReturn::StopIteration(_) => {
                 self.state
                     .stack
                     .extend(elements.into_iter().rev().map(Some));
-                return Ok(None);
-            }
-            core::cmp::Ordering::Greater => {
-                format!("too many values to unpack (expected {size})")
+                Ok(None)
             }
-            core::cmp::Ordering::Less => format!(
-                "not enough values to unpack (expected {}, got {})",
-                size,
-                elements.len()
-            ),
-        };
-        Err(vm.new_value_error(msg))
+        }
     }
 
     fn convert_value(
