diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 53cc26355b3..6c638f07eba 100644 --- a/Lib/test/test_itertools.py +++ b/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")) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 8f571a939ac..a8acdb11687 100644 --- a/vm/src/stdlib/itertools.rs +++ b/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, - fillvalue: PyObjectRef, + fillvalue: PyRwLock, } #[pyclass(with(IterNext, Constructor))] - impl PyItertoolsZipLongest {} + impl PyItertoolsZipLongest { + #[pymethod(magic)] + fn reduce(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + let args: Vec = 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, state: PyObjectRef, _vm: &VirtualMachine) -> PyResult<()> { + *zelf.fillvalue.write() = state; + Ok(()) + } + } impl IterNextIterable for PyItertoolsZipLongest {} impl IterNext for PyItertoolsZipLongest { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { @@ -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);