Fix some failing tests

ShaharNaveh · ShaharNaveh · commit 3d8cac3e8e55 · 2025-09-02T11:09:48.000+03:00
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
@@ -602,7 +602,7 @@ def test_compress(self):
                     next(testIntermediate)
                     self.assertEqual(list(op(testIntermediate)), list(result2))
 
-    @unittest.expectedFailure # TODO: RUSTPYTHON
+    @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: DeprecationWarning not triggered
     @pickle_deprecated
     def test_count(self):
         self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs
@@ -238,10 +238,10 @@ mod decl {
 
     #[derive(FromArgs)]
     struct CountNewArgs {
-        #[pyarg(positional, optional)]
+        #[pyarg(any, optional)]
         start: OptionalArg<PyObjectRef>,
 
-        #[pyarg(positional, optional)]
+        #[pyarg(any, optional)]
         step: OptionalArg<PyObjectRef>,
     }
 
@@ -1945,6 +1945,7 @@ mod decl {
         exhausted: AtomicCell<bool>,
         iterable: PyIter,
         n: AtomicCell<usize>,
+        strict: AtomicCell<bool>,
     }
 
     #[derive(FromArgs)]
@@ -1953,14 +1954,20 @@ mod decl {
         iterable_ref: PyObjectRef,
         #[pyarg(positional)]
         n: PyIntRef,
+        #[pyarg(named, default = false)]
+        strict: bool,
     }
 
     impl Constructor for PyItertoolsBatched {
         type Args = BatchedNewArgs;
 
         fn py_new(
             cls: PyTypeRef,
-            Self::Args { iterable_ref, n }: Self::Args,
+            Self::Args {
+                iterable_ref,
+                n,
+                strict,
+            }: Self::Args,
             vm: &VirtualMachine,
         ) -> PyResult {
             let n = n.as_bigint();
@@ -1976,6 +1983,7 @@ mod decl {
                 iterable,
                 n: AtomicCell::new(n),
                 exhausted: AtomicCell::new(false),
+                strict: AtomicCell::new(strict),
             }
             .into_ref_with_type(vm, cls)
             .map(Into::into)
@@ -2005,9 +2013,16 @@ mod decl {
                     }
                 }
             }
-            match result.len() {
+            let res_len = result.len();
+            match res_len {
                 0 => Ok(PyIterReturn::StopIteration(None)),
-                _ => Ok(PyIterReturn::Return(vm.ctx.new_tuple(result).into())),
+                _ => {
+                    if zelf.strict.load() && res_len != n {
+                        Err(vm.new_value_error("batched(): incomplete batch"))
+                    } else {
+                        Ok(PyIterReturn::Return(vm.ctx.new_tuple(result).into()))
+                    }
+                }
             }
         }
     }
