diff --git a/tests/snippets/stdlib_itertools.py b/tests/snippets/stdlib_itertools.py index 66342f871..67664ebc2 100644 --- a/tests/snippets/stdlib_itertools.py +++ b/tests/snippets/stdlib_itertools.py @@ -200,4 +200,13 @@ assert 0 == next(it) assert None == next(it) assert False == next(it) assert [] == next(it) -assert {} == next(it) \ No newline at end of file +assert {} == next(it) + + +# itertools.dropwhile +it = itertools.dropwhile(lambda x: x<5, [1,4,6,4,1]) +assert 6 == next(it) +assert 4 == next(it) +assert 1 == next(it) +with assertRaises(StopIteration): + next(it) \ No newline at end of file diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs index 85165466d..e353532fa 100644 --- a/vm/src/pyobject.rs +++ b/vm/src/pyobject.rs @@ -670,7 +670,7 @@ where } } -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct PyCallable { obj: PyObjectRef, } diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index e5ed280a3..1ba74f6d2 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -1,4 +1,4 @@ -use std::cell::RefCell; +use std::cell::{Cell, RefCell}; use std::cmp::Ordering; use std::ops::{AddAssign, SubAssign}; @@ -12,7 +12,7 @@ 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::{IdProtocol, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::pyobject::{IdProtocol, PyCallable, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue}; use crate::vm::VirtualMachine; #[pyclass(name = "chain")] @@ -290,6 +290,67 @@ impl PyItertoolsTakewhile { } } +#[pyclass] +#[derive(Debug)] +struct PyItertoolsDropwhile { + predicate: PyCallable, + iterable: PyObjectRef, + start_flag: Cell, +} + +impl PyValue for PyItertoolsDropwhile { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "dropwhile") + } +} + +type PyItertoolsDropwhileRef = PyRef; + +#[pyimpl] +impl PyItertoolsDropwhile { + #[pymethod(name = "__new__")] + #[allow(clippy::new_ret_no_self)] + fn new( + cls: PyClassRef, + predicate: PyCallable, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + let iter = get_iter(vm, &iterable)?; + + PyItertoolsDropwhile { + predicate, + iterable: iter, + start_flag: Cell::new(false), + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + let predicate = &self.predicate; + let iterable = &self.iterable; + + if !self.start_flag.get() { + loop { + let obj = call_next(vm, iterable)?; + let pred = predicate.clone(); + let pred_value = vm.invoke(&pred.into_object(), vec![obj.clone()])?; + if !objbool::boolval(vm, pred_value)? { + self.start_flag.set(true); + return Ok(obj); + } + } + } + call_next(vm, iterable) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { + zelf + } +} + #[pyclass(name = "islice")] #[derive(Debug)] struct PyItertoolsIslice { @@ -484,6 +545,9 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let count = ctx.new_class("count", ctx.object()); PyItertoolsCount::extend_class(ctx, &count); + let dropwhile = ctx.new_class("dropwhile", ctx.object()); + PyItertoolsDropwhile::extend_class(ctx, &dropwhile); + let repeat = ctx.new_class("repeat", ctx.object()); PyItertoolsRepeat::extend_class(ctx, &repeat); @@ -500,6 +564,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { py_module!(vm, "itertools", { "chain" => chain, "count" => count, + "dropwhile" => dropwhile, "repeat" => repeat, "starmap" => starmap, "takewhile" => takewhile,