Merge pull request #4390 from garychia/islice-reduce

Implement islice.__reduce__

extra_tests/snippets/stdlib_itertools.py

@@ -260,28 +260,39 @@ for proto in range(pickle.HIGHEST_PROTOCOL + 1):
 def assert_matches_seq(it, seq):
     assert list(it) == list(seq)
 
+def test_islice_pickle(it):
+    for p in range(pickle.HIGHEST_PROTOCOL + 1):
+        it == pickle.loads(pickle.dumps(it, p))
+
 i = itertools.islice
 
 it = i([1, 2, 3, 4, 5], 3)
 assert_matches_seq(it, [1, 2, 3])
+test_islice_pickle(it)
 
 it = i([0.5, 1, 1.5, 2, 2.5, 3, 4, 5], 1, 6, 2)
 assert_matches_seq(it, [1, 2, 3])
+test_islice_pickle(it)
 
 it = i([1, 2], None)
 assert_matches_seq(it, [1, 2])
+test_islice_pickle(it)
 
 it = i([1, 2, 3], None, None, None)
 assert_matches_seq(it, [1, 2, 3])
+test_islice_pickle(it)
 
 it = i([1, 2, 3], 1, None, None)
 assert_matches_seq(it, [2, 3])
+test_islice_pickle(it)
 
 it = i([1, 2, 3], None, 2, None)
 assert_matches_seq(it, [1, 2])
+test_islice_pickle(it)
 
 it = i([1, 2, 3], None, None, 3)
 assert_matches_seq(it, [1])
+test_islice_pickle(it)
 
 # itertools.filterfalse
 it = itertools.filterfalse(lambda x: x%2, range(10))

vm/src/stdlib/itertools.rs

@@ -924,6 +924,35 @@ mod decl {
             .into_ref_with_type(vm, cls)
             .map(Into::into)
         }
+
+        #[pymethod(magic)]
+        fn reduce(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
+            let cls = zelf.class().to_owned();
+            let itr = zelf.iterable.clone();
+            let cur = zelf.cur.take();
+            let next = zelf.next.take();
+            let step = zelf.step;
+            match zelf.stop {
+                Some(stop) => Ok(vm.new_tuple((cls, (itr, next, stop, step), (cur,)))),
+                _ => Ok(vm.new_tuple((cls, (itr, next, vm.new_pyobj(()), step), (cur,)))),
+            }
+        }
+
+        #[pymethod(magic)]
+        fn setstate(zelf: PyRef<Self>, state: PyTupleRef, vm: &VirtualMachine) -> PyResult<()> {
+            let args = state.as_slice();
+            if args.len() != 1 {
+                let msg = format!("function takes exactly 1 argument ({} given)", args.len());
+                return Err(vm.new_type_error(msg));
+            }
+            let cur = &args[0];
+            if let Ok(cur) = usize::try_from_object(vm, cur.clone()) {
+                zelf.cur.store(cur);
+            } else {
+                return Err(vm.new_type_error(String::from("Argument must be usize.")));
+            }
+            Ok(())
+        }
     }
 
     impl IterNextIterable for PyItertoolsIslice {}