Implement strict mode for zip() function to enforce equal lengths of iterables (#324)

Co-authored-by: amitrechavia <93474585+amitrechavia@users.noreply.github.com>

Author: davidhewitt

crates/monty/src/builtins/zip.rs

@@ -1,27 +1,27 @@
 //! Implementation of the zip() builtin function.
 
 use crate::{
-    args::ArgValues,
+    args::{ArgValues, KwargsValues},
     bytecode::VM,
-    defer_drop_mut,
-    exception_private::RunResult,
-    heap::HeapData,
+    defer_drop, defer_drop_mut,
+    exception_private::{ExcType, RunError, RunResult, SimpleException},
+    heap::{HeapData, HeapGuard},
     resource::ResourceTracker,
-    types::{List, MontyIter, allocate_tuple, tuple::TupleVec},
+    types::{List, MontyIter, PyTrait, allocate_tuple, tuple::TupleVec},
     value::Value,
 };
 
 /// Implementation of the zip() builtin function.
 ///
 /// Returns a list of tuples, where the i-th tuple contains the i-th element
 /// from each of the argument iterables. Stops when the shortest iterable is exhausted.
+/// When `strict=True`, raises `ValueError` if any iterable has a different length.
 /// Note: In Python this returns an iterator, but we return a list for simplicity.
 pub fn builtin_zip(vm: &mut VM<'_, '_, impl ResourceTracker>, args: ArgValues) -> RunResult<Value> {
     let (positional, kwargs) = args.into_parts();
     defer_drop_mut!(positional, vm);
 
-    // TODO: support kwargs (strict)
-    kwargs.not_supported_yet("zip", vm.heap)?;
+    let strict = extract_zip_strict(kwargs, vm)?;
 
     if positional.len() == 0 {
         // zip() with no arguments returns empty list
@@ -30,48 +30,113 @@ pub fn builtin_zip(vm: &mut VM<'_, '_, impl ResourceTracker>, args: ArgValues) -
     }
 
     // Create iterators for each iterable
-    let mut iterators: Vec<MontyIter> = Vec::with_capacity(positional.len());
+    let iterators: Vec<MontyIter> = Vec::with_capacity(positional.len());
+    defer_drop_mut!(iterators, vm);
     for iterable in positional {
-        match MontyIter::new(iterable, vm) {
-            Ok(iter) => iterators.push(iter),
-            Err(e) => {
-                // Clean up already-created iterators
-                for iter in iterators {
-                    iter.drop_with_heap(vm);
-                }
-                return Err(e);
-            }
-        }
+        iterators.push(MontyIter::new(iterable, vm)?);
     }
 
-    let mut result: Vec<Value> = Vec::new();
+    let mut result_guard = HeapGuard::new(Vec::new(), vm);
+    let (result, vm) = result_guard.as_parts_mut();
 
     // Zip until shortest iterator is exhausted
     'outer: loop {
-        let mut tuple_items = TupleVec::with_capacity(iterators.len());
+        let mut items_guard = HeapGuard::new(TupleVec::with_capacity(iterators.len()), vm);
+        let (tuple_items, vm) = items_guard.as_parts_mut();
 
-        for iter in &mut iterators {
+        for (i, iter) in iterators.iter_mut().enumerate() {
             if let Some(item) = iter.for_next(vm)? {
                 tuple_items.push(item);
             } else {
-                // This iterator is exhausted - drop partial tuple items and stop
-                for item in tuple_items {
-                    item.drop_with_heap(vm);
+                // This iterator is exhausted - stop zipping
+
+                if strict {
+                    // In strict mode, if i > 0 then argument i+1 ran out before
+                    // the earlier ones, so it is "shorter."
+                    if i > 0 {
+                        return Err(strict_length_error(i + 1, i, "shorter"));
+                    }
+                    // i == 0: first iterator exhausted — verify every remaining
+                    // iterator is also exhausted. If any still yields a value,
+                    // that argument is "longer" than all preceding exhausted ones.
+                    // j is the 0-based index; iterators 0..j are all exhausted,
+                    // so j gives the count for the error message.
+                    for (j, remaining) in iterators.iter_mut().enumerate().skip(1) {
+                        if let Some(extra) = remaining.for_next(vm)? {
+                            extra.drop_with_heap(vm);
+                            return Err(strict_length_error(j + 1, j, "longer"));
+                        }
+                    }
                 }
+
                 break 'outer;
             }
         }
 
         // Create tuple from collected items
+        let (tuple_items, vm) = items_guard.into_parts();
         let tuple_val = allocate_tuple(tuple_items, vm.heap)?;
         result.push(tuple_val);
     }
 
-    // Clean up iterators
-    for iter in iterators {
-        iter.drop_with_heap(vm);
-    }
-
+    let (result, vm) = result_guard.into_parts();
     let heap_id = vm.heap.allocate(HeapData::List(List::new(result)))?;
     Ok(Value::Ref(heap_id))
 }
+
+/// Extracts the `strict` keyword argument from `zip()`.
+///
+/// Accepts any truthy/falsy value for `strict`, matching CPython behavior.
+/// Raises `TypeError` for unexpected keyword arguments.
+fn extract_zip_strict(kwargs: KwargsValues, vm: &mut VM<'_, '_, impl ResourceTracker>) -> RunResult<bool> {
+    let mut strict = false;
+    let mut error: Option<RunError> = None;
+
+    for (key, value) in kwargs {
+        defer_drop!(key, vm);
+        defer_drop!(value, vm);
+
+        if error.is_some() {
+            continue;
+        }
+
+        let Some(keyword_name) = key.as_either_str(vm.heap) else {
+            error = Some(SimpleException::new_msg(ExcType::TypeError, "keywords must be strings").into());
+            continue;
+        };
+
+        let key_str = keyword_name.as_str(vm.interns);
+        match key_str {
+            "strict" => {
+                strict = value.py_bool(vm);
+            }
+            _ => {
+                error = Some(ExcType::type_error_unexpected_keyword("zip", key_str));
+            }
+        }
+    }
+
+    if let Some(error) = error {
+        Err(error)
+    } else {
+        Ok(strict)
+    }
+}
+
+/// Builds the `ValueError` for `zip(strict=True)` when iterables have different lengths.
+///
+/// Matches CPython's error format:
+/// - `"zip() argument 2 is shorter than argument 1"` (singular)
+/// - `"zip() argument 4 is shorter than arguments 1-3"` (plural)
+fn strict_length_error(exhausted_arg: usize, num_longer_args: usize, relation: &str) -> RunError {
+    let others = if num_longer_args == 1 {
+        "argument 1".to_owned()
+    } else {
+        format!("arguments 1-{num_longer_args}")
+    };
+    SimpleException::new_msg(
+        ExcType::ValueError,
+        format!("zip() argument {exhausted_arg} is {relation} than {others}"),
+    )
+    .into()
+}

crates/monty/src/heap_traits.rs

@@ -4,6 +4,8 @@ use std::{
     vec::{Drain, IntoIter},
 };
 
+use smallvec::SmallVec;
+
 use crate::{
     ResourceTracker,
     heap::{Heap, HeapId, RecursionToken},
@@ -112,6 +114,17 @@ impl<U: DropWithHeap> DropWithHeap for Vec<U> {
     }
 }
 
+impl<A: smallvec::Array> DropWithHeap for SmallVec<A>
+where
+    A::Item: DropWithHeap,
+{
+    fn drop_with_heap<H: ContainsHeap>(self, heap: &mut H) {
+        for value in self {
+            value.drop_with_heap(heap);
+        }
+    }
+}
+
 impl<U: DropWithHeap> DropWithHeap for IntoIter<U> {
     fn drop_with_heap<H: ContainsHeap>(self, heap: &mut H) {
         for value in self {

crates/monty/test_cases/builtin__more_iter_funcs.py

@@ -359,3 +359,59 @@ def negate(x):
 # zip with empty
 assert list(zip([1, 2], [])) == [], 'zip with empty second'
 assert list(zip([], [1, 2])) == [], 'zip with empty first'
+
+# === zip(strict=True) ===
+# Equal length iterables succeed
+assert list(zip([1, 2], [3, 4], strict=True)) == [(1, 3), (2, 4)], 'zip strict equal lengths'
+assert list(zip([1], [2], [3], strict=True)) == [(1, 2, 3)], 'zip strict three single-element lists'
+assert list(zip([], [], strict=True)) == [], 'zip strict empty lists'
+assert list(zip(strict=True)) == [], 'zip strict no arguments'
+assert list(zip([1, 2, 3], strict=True)) == [(1,), (2,), (3,)], 'zip strict single iterable'
+
+# strict=False behaves like default
+assert list(zip([1, 2, 3], [4, 5], strict=False)) == [(1, 4), (2, 5)], 'zip strict=False truncates'
+
+# Falsy values are accepted
+assert list(zip([1, 2, 3], [4, 5], strict=0)) == [(1, 4), (2, 5)], 'zip strict=0 is falsy'
+
+# Second argument shorter
+try:
+    list(zip([1, 2, 3], [4, 5], strict=True))
+    assert False, 'zip strict should raise for shorter arg 2'
+except ValueError as e:
+    assert str(e) == 'zip() argument 2 is shorter than argument 1', 'zip strict shorter error'
+
+# Second argument longer
+try:
+    list(zip([1, 2], [4, 5, 6], strict=True))
+    assert False, 'zip strict should raise for longer arg 2'
+except ValueError as e:
+    assert str(e) == 'zip() argument 2 is longer than argument 1', 'zip strict longer error'
+
+# Third argument shorter with plural
+try:
+    list(zip([1, 2], [3, 4], [5], strict=True))
+    assert False, 'zip strict should raise for shorter arg 3'
+except ValueError as e:
+    assert str(e) == 'zip() argument 3 is shorter than arguments 1-2', 'zip strict shorter plural'
+
+# Fourth argument shorter
+try:
+    list(zip([1, 2], [3, 4], [5, 6], [7], strict=True))
+    assert False, 'zip strict should raise for shorter arg 4'
+except ValueError as e:
+    assert str(e) == 'zip() argument 4 is shorter than arguments 1-3', 'zip strict shorter 4 args'
+
+# Third argument longer than arguments 1-2 (both exhausted)
+try:
+    list(zip([1], [2], [3, 4], strict=True))
+    assert False, 'zip strict should raise for longer arg 3'
+except ValueError as e:
+    assert str(e) == 'zip() argument 3 is longer than arguments 1-2', 'zip strict longer plural'
+
+# Unexpected keyword argument
+try:
+    list(zip([1], foo=True))
+    assert False, 'zip unexpected kwarg should raise TypeError'
+except TypeError as e:
+    assert str(e) == "zip() got an unexpected keyword argument 'foo'", 'zip unexpected kwarg error'