Merge pull request RustPython#3791 from rng-dynamics/feat-itertools-compress-reduce

Add itertools.compress.__reduce__

Author: youknowone

Lib/test/test_itertools.py

@@ -500,8 +500,6 @@ def test_combinatorics(self):
                 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr)))     # comb: cwr that is a perm
                 self.assertEqual(comb, sorted(set(cwr) & set(perm)))            # comb: both a cwr and a perm
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_compress(self):
         self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
         self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))

vm/src/stdlib/itertools.rs

@@ -116,39 +116,47 @@ mod decl {
     #[derive(Debug, PyPayload)]
     struct PyItertoolsCompress {
         data: PyIter,
-        selector: PyIter,
+        selectors: PyIter,
     }
 
     #[derive(FromArgs)]
     struct CompressNewArgs {
-        #[pyarg(positional)]
+        #[pyarg(any)]
         data: PyIter,
-        #[pyarg(positional)]
-        selector: PyIter,
+        #[pyarg(any)]
+        selectors: PyIter,
     }
 
     impl Constructor for PyItertoolsCompress {
         type Args = CompressNewArgs;
 
         fn py_new(
             cls: PyTypeRef,
-            Self::Args { data, selector }: Self::Args,
+            Self::Args { data, selectors }: Self::Args,
             vm: &VirtualMachine,
         ) -> PyResult {
-            PyItertoolsCompress { data, selector }
+            PyItertoolsCompress { data, selectors }
                 .into_ref_with_type(vm, cls)
                 .map(Into::into)
         }
     }
 
     #[pyimpl(with(IterNext, Constructor))]
-    impl PyItertoolsCompress {}
+    impl PyItertoolsCompress {
+        #[pymethod(magic)]
+        fn reduce(zelf: PyRef<Self>) -> (PyTypeRef, (PyIter, PyIter)) {
+            (
+                zelf.class().clone(),
+                (zelf.data.clone(), zelf.selectors.clone()),
+            )
+        }
+    }
 
     impl IterNextIterable for PyItertoolsCompress {}
     impl IterNext for PyItertoolsCompress {
         fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
             loop {
-                let sel_obj = match zelf.selector.next(vm)? {
+                let sel_obj = match zelf.selectors.next(vm)? {
                     PyIterReturn::Return(obj) => obj,
                     PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
                 };