Fix #462

ml31415 · ml31415 · commit 0edfa4052aa7 · 2022-12-21T16:09:29.000+01:00
diff --git a/streamz/core.py b/streamz/core.py
@@ -1,6 +1,7 @@
 import asyncio
 from collections import deque, defaultdict
 from datetime import timedelta
+from itertools import chain
 import functools
 import logging
 import six
@@ -1637,15 +1638,20 @@ class flatten(Stream):
     """
     def update(self, x, who=None, metadata=None):
         L = []
-        for i, item in enumerate(x):
-            if i == len(x) - 1:
-                y = self._emit(item, metadata=metadata)
-            else:
-                y = self._emit(item)
+        items = chain(x)
+        item = next(items)
+        for item_next in items:
+            y = self._emit(item)
+            item = item_next
             if type(y) is list:
                 L.extend(y)
             else:
                 L.append(y)
+        y = self._emit(item, metadata=metadata)
+        if type(y) is list:
+            L.extend(y)
+        else:
+            L.append(y)
         return L
 
 
diff --git a/streamz/tests/test_core.py b/streamz/tests/test_core.py
@@ -800,13 +800,15 @@ def test_frequencies():
     assert L[-1] == {'a': 2, 'b': 1}
 
 
-def test_flatten():
+@pytest.mark.parametrize("iterators",
+                         [[[1, 2, 3], [4, 5], [6, 7, 8]],
+                          [(i for i in range(1, 7)), (i for i in range(7, 9))]])
+def test_flatten(iterators):
     source = Stream()
     L = source.flatten().sink_to_list()
 
-    source.emit([1, 2, 3])
-    source.emit([4, 5])
-    source.emit([6, 7, 8])
+    for iterator in iterators:
+        source.emit(iterator)
 
     assert L == [1, 2, 3, 4, 5, 6, 7, 8]
 
