Merge pull request #965 from Jongy/itertools-takewhile

Itertools takewhile

vm/src/stdlib/itertools.rs

@@ -5,8 +5,9 @@ use std::ops::{AddAssign, SubAssign};
 use num_bigint::BigInt;
 
 use crate::function::OptionalArg;
+use crate::obj::objbool;
 use crate::obj::objint::{PyInt, PyIntRef};
-use crate::obj::objiter::new_stop_iteration;
+use crate::obj::objiter::{call_next, get_iter, new_stop_iteration};
 use crate::obj::objtype::PyClassRef;
 use crate::pyobject::{PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
@@ -121,6 +122,65 @@ impl PyItertoolsRepeat {
     }
 }
 
+#[pyclass]
+#[derive(Debug)]
+struct PyItertoolsTakewhile {
+    predicate: PyObjectRef,
+    iterable: PyObjectRef,
+    stop_flag: RefCell<bool>,
+}
+
+impl PyValue for PyItertoolsTakewhile {
+    fn class(vm: &VirtualMachine) -> PyClassRef {
+        vm.class("itertools", "takewhile")
+    }
+}
+
+#[pyimpl]
+impl PyItertoolsTakewhile {
+    #[pymethod(name = "__new__")]
+    fn new(
+        _cls: PyClassRef,
+        predicate: PyObjectRef,
+        iterable: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        let iter = get_iter(vm, &iterable)?;
+
+        Ok(PyItertoolsTakewhile {
+            predicate: predicate,
+            iterable: iter,
+            stop_flag: RefCell::new(false),
+        }
+        .into_ref(vm)
+        .into_object())
+    }
+
+    #[pymethod(name = "__next__")]
+    fn next(&self, vm: &VirtualMachine) -> PyResult {
+        if *self.stop_flag.borrow() {
+            return Err(new_stop_iteration(vm));
+        }
+
+        // might be StopIteration or anything else, which is propaged upwwards
+        let obj = call_next(vm, &self.iterable)?;
+
+        let verdict = vm.invoke(self.predicate.clone(), vec![obj.clone()])?;
+        let verdict = objbool::boolval(vm, verdict)?;
+        if verdict {
+            Ok(obj)
+        } else {
+            *self.stop_flag.borrow_mut() = true;
+            Err(new_stop_iteration(vm))
+        }
+    }
+
+    #[pymethod(name = "__iter__")]
+    fn iter(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
+        zelf
+    }
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
 
@@ -130,8 +190,12 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let repeat = ctx.new_class("repeat", ctx.object());
     PyItertoolsRepeat::extend_class(ctx, &repeat);
 
+    let takewhile = ctx.new_class("takewhile", ctx.object());
+    PyItertoolsTakewhile::extend_class(ctx, &takewhile);
+
     py_module!(vm, "itertools", {
         "count" => count,
         "repeat" => repeat,
+        "takewhile" => takewhile,
     })
 }