From 2c061c095d05eb984ea86ec821e5f4737017a4d6 Mon Sep 17 00:00:00 2001 From: Chia-Hsiang Cheng Date: Sat, 31 Dec 2022 21:23:10 +0800 Subject: [PATCH 1/2] Implement reduce for islice --- vm/src/stdlib/itertools.rs | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 2b605b920f0..b5f5686c42d 100644 --- a/vm/src/stdlib/itertools.rs +++ b/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, vm: &VirtualMachine) -> PyResult { + 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, 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 {} From 49597bc523c664b9d13b49add25f0b12d4925e53 Mon Sep 17 00:00:00 2001 From: Chia-Hsiang Cheng Date: Sat, 31 Dec 2022 21:23:37 +0800 Subject: [PATCH 2/2] Add tests for islice --- extra_tests/snippets/stdlib_itertools.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/extra_tests/snippets/stdlib_itertools.py b/extra_tests/snippets/stdlib_itertools.py index 0fe51bb82f4..58684f611d8 100644 --- a/extra_tests/snippets/stdlib_itertools.py +++ b/extra_tests/snippets/stdlib_itertools.py @@ -260,28 +260,39 @@ def underten(x): 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))