From e1854cec25311eb0520dfa5b807544a5f4277d59 Mon Sep 17 00:00:00 2001 From: Dong-hee Na Date: Wed, 28 Aug 2019 12:27:52 +0900 Subject: [PATCH 1/2] itertools.dropwhile: Implement itertools.dropwhile --- tests/snippets/stdlib_itertools.py | 11 ++++- vm/src/stdlib/itertools.rs | 67 ++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 1 deletion(-) 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/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index e5ed280a3..5e5cd6afd 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -290,6 +290,69 @@ impl PyItertoolsTakewhile { } } +#[pyclass] +#[derive(Debug)] +struct PyItertoolsDropwhile { + predicate: PyObjectRef, + iterable: PyObjectRef, + start_flag: RefCell, +} + +impl PyValue for PyItertoolsDropwhile { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "dropwhile") + } +} + +#[pyimpl] +impl PyItertoolsDropwhile { + #[pymethod(name = "__new__")] + #[allow(clippy::new_ret_no_self)] + fn new( + _cls: PyClassRef, + predicate: PyObjectRef, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + let iter = get_iter(vm, &iterable)?; + + Ok(PyItertoolsDropwhile { + predicate, + iterable: iter, + start_flag: RefCell::new(false), + } + .into_ref(vm) + .into_object()) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + let predicate = &self.predicate; + let iterable = &self.iterable; + + if !*self.start_flag.borrow_mut() { + loop { + let obj = call_next(vm, iterable)?; + let pred_value = vm.invoke(predicate, vec![obj.clone()])?; + if !objbool::boolval(vm, pred_value)? { + *self.start_flag.borrow_mut() = true; + return Ok(obj); + } + } + } + + loop { + let obj = call_next(vm, iterable)?; + return Ok(obj); + } + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { + zelf + } +} + #[pyclass(name = "islice")] #[derive(Debug)] struct PyItertoolsIslice { @@ -484,6 +547,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 +566,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { py_module!(vm, "itertools", { "chain" => chain, "count" => count, + "dropwhile" => dropwhile, "repeat" => repeat, "starmap" => starmap, "takewhile" => takewhile, From 44ceaa8dac612085c074419e0c93947af11face0 Mon Sep 17 00:00:00 2001 From: Dong-hee Na Date: Wed, 28 Aug 2019 18:51:41 +0900 Subject: [PATCH 2/2] itertools: Apply codereview --- vm/src/pyobject.rs | 2 +- vm/src/stdlib/itertools.rs | 36 +++++++++++++++++------------------- 2 files changed, 18 insertions(+), 20 deletions(-) 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 5e5cd6afd..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")] @@ -293,9 +293,9 @@ impl PyItertoolsTakewhile { #[pyclass] #[derive(Debug)] struct PyItertoolsDropwhile { - predicate: PyObjectRef, + predicate: PyCallable, iterable: PyObjectRef, - start_flag: RefCell, + start_flag: Cell, } impl PyValue for PyItertoolsDropwhile { @@ -304,25 +304,26 @@ impl PyValue for PyItertoolsDropwhile { } } +type PyItertoolsDropwhileRef = PyRef; + #[pyimpl] impl PyItertoolsDropwhile { #[pymethod(name = "__new__")] #[allow(clippy::new_ret_no_self)] fn new( - _cls: PyClassRef, - predicate: PyObjectRef, + cls: PyClassRef, + predicate: PyCallable, iterable: PyObjectRef, vm: &VirtualMachine, - ) -> PyResult { + ) -> PyResult { let iter = get_iter(vm, &iterable)?; - Ok(PyItertoolsDropwhile { + PyItertoolsDropwhile { predicate, iterable: iter, - start_flag: RefCell::new(false), + start_flag: Cell::new(false), } - .into_ref(vm) - .into_object()) + .into_ref_with_type(vm, cls) } #[pymethod(name = "__next__")] @@ -330,21 +331,18 @@ impl PyItertoolsDropwhile { let predicate = &self.predicate; let iterable = &self.iterable; - if !*self.start_flag.borrow_mut() { + if !self.start_flag.get() { loop { let obj = call_next(vm, iterable)?; - let pred_value = vm.invoke(predicate, vec![obj.clone()])?; + let pred = predicate.clone(); + let pred_value = vm.invoke(&pred.into_object(), vec![obj.clone()])?; if !objbool::boolval(vm, pred_value)? { - *self.start_flag.borrow_mut() = true; + self.start_flag.set(true); return Ok(obj); } } } - - loop { - let obj = call_next(vm, iterable)?; - return Ok(obj); - } + call_next(vm, iterable) } #[pymethod(name = "__iter__")]