Ensure flatten.update([]) does not kill the stream

Trap the StopIteration that erupts from update when passed an empty
iterable so that the stream can continue.

Author: dwskoog

examples/scrape.py

@@ -5,7 +5,7 @@
 import toolz
 from bs4 import BeautifulSoup
 
-from streamz import Stream, identity
+from streamz import Stream
 
 
 def links_of_page(content_page):
@@ -38,7 +38,6 @@ def topk_dict(d, k=10):
                 .map(lambda x: x.content))
 links = (content.zip(pages)
                 .map(links_of_page)
-                .filter(identity)
                 .flatten())
 links.connect(source)
 

streamz/core.py

@@ -1637,7 +1637,10 @@ class flatten(Stream):
     def update(self, x, who=None, metadata=None):
         L = []
         items = chain(x)
-        item = next(items)
+        try:
+            item = next(items)
+        except StopIteration:
+            return L
         for item_next in items:
             y = self._emit(item)
             item = item_next

streamz/tests/test_core.py

@@ -814,6 +814,17 @@ def test_flatten(iterators):
     assert L == [1, 2, 3, 4, 5, 6, 7, 8]
 
 
+def test_flatten_empty():
+    source = Stream()
+    L = source.flatten().sink_to_list()
+
+    source.emit([1, 2])
+    source.emit([])
+    source.emit([3, 4])
+
+    assert L == [1, 2, 3, 4]
+
+
 def test_unique():
     source = Stream()
     L = source.unique().sink_to_list()