Implement itertools.chain.from_iterable()

dralley · dralley · commit c1b5aa719d71 · 2019-12-24T22:53:11.000-05:00
diff --git a/tests/snippets/stdlib_itertools.py b/tests/snippets/stdlib_itertools.py
@@ -21,6 +21,36 @@
 with assert_raises(TypeError):
     next(x)
 
+# empty
+with assert_raises(TypeError):
+    chain.from_iterable()
+
+with assert_raises(TypeError):
+    chain.from_iterable("abc", "def")
+
+with assert_raises(TypeError):
+    # iterables are lazily evaluated -- can be constructed but will fail to execute
+    list(chain.from_iterable([1, 2, 3]))
+
+with assert_raises(TypeError):
+    list(chain(1))
+
+args = ["abc", "def"]
+assert list(chain.from_iterable(args)) == ['a', 'b', 'c', 'd', 'e', 'f']
+
+args = [[], "", b"", ()]
+assert list(chain.from_iterable(args)) == []
+
+args = ["ab", "cd", (), 'e']
+assert list(chain.from_iterable(args)) == ['a', 'b', 'c', 'd', 'e']
+
+x = chain.from_iterable(["ab", 1])
+assert next(x) == 'a'
+assert next(x) == 'b'
+with assert_raises(TypeError):
+    next(x)
+
+
 # itertools.count tests
 
 # default arguments
@@ -117,7 +147,7 @@
 with assert_raises(StopIteration):
     next(r)
 
-# timees = 0
+# times = 0
 r = itertools.repeat(1, 0)
 with assert_raises(StopIteration):
     next(r)
diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs
@@ -19,6 +19,7 @@ use crate::pyobject::{
 };
 use crate::vm::VirtualMachine;
 
+
 #[pyclass(name = "chain")]
 #[derive(Debug)]
 struct PyItertoolsChain {
@@ -73,6 +74,23 @@ impl PyItertoolsChain {
     fn iter(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
         zelf
     }
+
+    #[pyclassmethod(name = "from_iterable")]
+    fn from_iterable(
+        cls: PyClassRef,
+        iterable: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult<PyRef<Self>> {
+
+        let it = get_iter(vm, &iterable)?;
+        let iterables = get_all(vm, &it)?;
+
+        PyItertoolsChain {
+            iterables: iterables,
+            cur: RefCell::new((0, None)),
+        }
+        .into_ref_with_type(vm, cls)
+    }
 }
 
 #[pyclass(name = "compress")]
