Merge pull request RustPython#4235 from dannasman/zip_longest_reduce

Add reduce and setstate for zip_longest

Author: youknowone

Lib/test/test_itertools.py

@@ -985,8 +985,6 @@ def test_zip_longest_tuple_reuse(self):
         ids = list(map(id, list(zip_longest('abc', 'def'))))
         self.assertEqual(len(dict.fromkeys(ids)), len(ids))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_zip_longest_pickling(self):
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             self.pickletest(proto, zip_longest("abc", "def"))

vm/src/stdlib/itertools.rs

@@ -1593,7 +1593,7 @@ mod decl {
             let iterators = iterators.into_vec();
             PyItertoolsZipLongest {
                 iterators,
-                fillvalue,
+                fillvalue: PyRwLock::new(fillvalue),
             }
             .into_ref_with_type(vm, cls)
             .map(Into::into)
@@ -1605,11 +1605,31 @@ mod decl {
     #[derive(Debug, PyPayload)]
     struct PyItertoolsZipLongest {
         iterators: Vec<PyIter>,
-        fillvalue: PyObjectRef,
+        fillvalue: PyRwLock<PyObjectRef>,
     }
 
     #[pyclass(with(IterNext, Constructor))]
-    impl PyItertoolsZipLongest {}
+    impl PyItertoolsZipLongest {
+        #[pymethod(magic)]
+        fn reduce(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
+            let args: Vec<PyObjectRef> = zelf
+                .iterators
+                .iter()
+                .map(|i| i.clone().to_pyobject(vm))
+                .collect();
+            Ok(vm.new_tuple((
+                zelf.class().to_owned(),
+                vm.new_tuple(args),
+                zelf.fillvalue.read().to_owned(),
+            )))
+        }
+
+        #[pymethod(magic)]
+        fn setstate(zelf: PyRef<Self>, state: PyObjectRef, _vm: &VirtualMachine) -> PyResult<()> {
+            *zelf.fillvalue.write() = state;
+            Ok(())
+        }
+    }
     impl IterNextIterable for PyItertoolsZipLongest {}
     impl IterNext for PyItertoolsZipLongest {
         fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
@@ -1627,7 +1647,7 @@ mod decl {
                         if numactive == 0 {
                             return Ok(PyIterReturn::StopIteration(v));
                         }
-                        zelf.fillvalue.clone()
+                        zelf.fillvalue.read().clone()
                     }
                 };
                 result.push(next_obj);