diff --git a/tests/snippets/stdlib_itertools.py b/tests/snippets/stdlib_itertools.py index c8921cb91..e8285fac9 100644 --- a/tests/snippets/stdlib_itertools.py +++ b/tests/snippets/stdlib_itertools.py @@ -154,3 +154,32 @@ with assertRaises(StopIteration): next(t) with assertRaises(StopIteration): next(t) + + +# itertools.islice tests + +def assert_matches_seq(it, seq): + assert list(it) == list(seq) + +i = itertools.islice + +it = i([1, 2, 3, 4, 5], 3) +assert_matches_seq(it, [1, 2, 3]) + +it = i([0.5, 1, 1.5, 2, 2.5, 3, 4, 5], 1, 6, 2) +assert_matches_seq(it, [1, 2, 3]) + +it = i([1, 2], None) +assert_matches_seq(it, [1, 2]) + +it = i([1, 2, 3], None, None, None) +assert_matches_seq(it, [1, 2, 3]) + +it = i([1, 2, 3], 1, None, None) +assert_matches_seq(it, [2, 3]) + +it = i([1, 2, 3], None, 2, None) +assert_matches_seq(it, [1, 2]) + +it = i([1, 2, 3], None, None, 3) +assert_matches_seq(it, [1]) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 7380a7677..a8c333701 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -3,14 +3,16 @@ use std::cmp::Ordering; use std::ops::{AddAssign, SubAssign}; use num_bigint::BigInt; +use num_traits::ToPrimitive; use crate::function::{OptionalArg, PyFuncArgs}; use crate::obj::objbool; +use crate::obj::objint; use crate::obj::objint::{PyInt, PyIntRef}; use crate::obj::objiter::{call_next, get_iter, new_stop_iteration}; use crate::obj::objtype; use crate::obj::objtype::PyClassRef; -use crate::pyobject::{PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::pyobject::{IdProtocol, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue}; use crate::vm::VirtualMachine; #[pyclass(name = "chain")] @@ -284,6 +286,133 @@ impl PyItertoolsTakewhile { } } +#[pyclass(name = "islice")] +#[derive(Debug)] +struct PyItertoolsIslice { + iterable: PyObjectRef, + cur: RefCell, + next: RefCell, + stop: Option, + step: usize, +} + +impl PyValue for PyItertoolsIslice { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "islice") + } +} + +fn pyobject_to_opt_usize(obj: PyObjectRef, vm: &VirtualMachine) -> Option { + let is_int = objtype::isinstance(&obj, &vm.ctx.int_type()); + if is_int { + objint::get_value(&obj).to_usize() + } else { + None + } +} + +#[pyimpl] +impl PyItertoolsIslice { + #[pymethod(name = "__new__")] + fn new(_cls: PyClassRef, args: PyFuncArgs, vm: &VirtualMachine) -> PyResult { + let (iter, start, stop, step) = match args.args.len() { + 0 | 1 => { + return Err(vm.new_type_error(format!( + "islice expected at least 2 arguments, got {}", + args.args.len() + ))); + } + + 2 => { + let (iter, stop): (PyObjectRef, PyObjectRef) = args.bind(vm)?; + + (iter, 0usize, stop, 1usize) + } + _ => { + let (iter, start, stop, step): ( + PyObjectRef, + PyObjectRef, + PyObjectRef, + PyObjectRef, + ) = args.bind(vm)?; + + let start = if !start.is(&vm.get_none()) { + pyobject_to_opt_usize(start, &vm).ok_or_else(|| { + vm.new_value_error( + "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.".to_string(), + ) + })? + } else { + 0usize + }; + + let step = if !step.is(&vm.get_none()) { + pyobject_to_opt_usize(step, &vm).ok_or_else(|| { + vm.new_value_error( + "Step for islice() must be a positive integer or None.".to_string(), + ) + })? + } else { + 1usize + }; + + (iter, start, stop, step) + } + }; + + let stop = if !stop.is(&vm.get_none()) { + Some(pyobject_to_opt_usize(stop, &vm).ok_or_else(|| { + vm.new_value_error( + "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize." + .to_string(), + ) + })?) + } else { + None + }; + + let iter = get_iter(vm, &iter)?; + + Ok(PyItertoolsIslice { + iterable: iter, + cur: RefCell::new(0), + next: RefCell::new(start), + stop: stop, + step: step, + } + .into_ref(vm) + .into_object()) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + while *self.cur.borrow() < *self.next.borrow() { + call_next(vm, &self.iterable)?; + *self.cur.borrow_mut() += 1; + } + + if let Some(stop) = self.stop { + if *self.cur.borrow() >= stop { + return Err(new_stop_iteration(vm)); + } + } + + let obj = call_next(vm, &self.iterable)?; + *self.cur.borrow_mut() += 1; + + // TODO is this overflow check required? attempts to copy CPython. + let (next, ovf) = (*self.next.borrow()).overflowing_add(self.step); + *self.next.borrow_mut() = if ovf { self.stop.unwrap() } else { next }; + + Ok(obj) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { + zelf + } +} + pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let ctx = &vm.ctx; @@ -300,11 +429,14 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let takewhile = ctx.new_class("takewhile", ctx.object()); PyItertoolsTakewhile::extend_class(ctx, &takewhile); + let islice = PyItertoolsIslice::make_class(ctx); + py_module!(vm, "itertools", { "chain" => chain, "count" => count, "repeat" => repeat, "starmap" => starmap, "takewhile" => takewhile, + "islice" => islice, }) }