itertools.compress

Hyunji Kim · Hyunji Kim · commit 1c8224df5a04 · 2019-09-22T21:38:07.000+09:00
diff --git a/tests/snippets/stdlib_itertools.py b/tests/snippets/stdlib_itertools.py
@@ -235,3 +235,11 @@ def assert_matches_seq(it, seq):
 assert 0 == next(it)
 with assert_raises(StopIteration):
     next(it)
+
+# itertools.compress
+assert list(itertools.compress("ABCDEF", [1,0,1,0,1,1])) == list("ACEF")
+assert list(itertools.compress("ABCDEF", [0,0,0,0,0,0])) == list("")
+assert list(itertools.compress("ABCDEF", [1,1,1,1,1,1])) == list("ABCDEF")
+assert list(itertools.compress("ABCDEF", [1,0,1])) == list("AC")
+assert list(itertools.compress("ABC", [0,1,1,1,1,1])) == list("BC")
+assert list(itertools.compress("ABCDEF", [True,False,"t","",1,9])) == list("ACEF")
diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs
@@ -73,6 +73,63 @@ impl PyItertoolsChain {
     }
 }
 
+#[pyclass(name = "compress")]
+#[derive(Debug)]
+struct PyItertoolsCompress {
+    data: PyObjectRef,
+    selector: PyObjectRef,
+}
+
+impl PyValue for PyItertoolsCompress {
+    fn class(vm: &VirtualMachine) -> PyClassRef {
+        vm.class("itertools", "compress")
+    }
+}
+
+#[pyimpl]
+impl PyItertoolsCompress {
+    #[pymethod(name = "__new__")]
+    #[allow(clippy::new_ret_no_self)]
+    fn new(
+        _cls: PyClassRef,
+        data: PyObjectRef,
+        selector: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        let data_iter = get_iter(vm, &data)?;
+        let selector_iter = get_iter(vm, &selector)?;
+
+        Ok(PyItertoolsCompress {
+            data: data_iter,
+            selector: selector_iter,
+        }
+        .into_ref(vm)
+        .into_object())
+    }
+
+    #[pymethod(name = "__next__")]
+    fn next(&self, vm: &VirtualMachine) -> PyResult {
+        loop {
+            let selector_iter = get_iter(vm, &self.selector)?;
+            let sel_obj = call_next(vm, &selector_iter)?;
+
+            let verdict = objbool::boolval(vm, sel_obj.clone())?;
+
+            let data_iter = get_iter(vm, &self.data)?;
+            let data_obj = call_next(vm, &data_iter)?;
+
+            if verdict {
+                return Ok(data_obj);
+            }
+        }
+    }
+
+    #[pymethod(name = "__iter__")]
+    fn iter(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
+        zelf
+    }
+}
+
 #[pyclass]
 #[derive(Debug)]
 struct PyItertoolsCount {
@@ -577,8 +634,8 @@ impl PyItertoolsAccumulate {
         let obj = call_next(vm, iterable)?;
 
         let next_acc_value = match &*self.acc_value.borrow() {
-            Option::None => obj.clone(),
-            Option::Some(value) => {
+            None => obj.clone(),
+            Some(value) => {
                 if self.binop.is(&vm.get_none()) {
                     vm._add(value.clone(), obj.clone())?
                 } else {
@@ -602,6 +659,8 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
 
     let chain = PyItertoolsChain::make_class(ctx);
 
+    let compress = PyItertoolsCompress::make_class(ctx);
+
     let count = ctx.new_class("count", ctx.object());
     PyItertoolsCount::extend_class(ctx, &count);
 
@@ -626,6 +685,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
 
     py_module!(vm, "itertools", {
         "chain" => chain,
+        "compress" => compress,
         "count" => count,
         "dropwhile" => dropwhile,
         "repeat" => repeat,
