Merge pull request #393 from roveo/from_iterable2

from_iterable source

Author: martindurant

docs/source/api.rst

@@ -51,6 +51,7 @@ Sources
 -------
 
 .. autosummary::
+   from_iterable
    filenames
    from_kafka
    from_kafka_batched
@@ -95,6 +96,7 @@ Definitions
 .. autofunction:: zip
 .. autofunction:: zip_latest
 
+.. autofunction:: from_iterable
 .. autofunction:: filenames
 .. autofunction:: from_kafka
 .. autofunction:: from_kafka_batched

streamz/sinks.py

@@ -17,6 +17,10 @@ def __init__(self, upstream, **kwargs):
         super().__init__(upstream, **kwargs)
         _global_sinks.add(self)
 
+    def destroy(self):
+        super().destroy()
+        _global_sinks.remove(self)
+
 
 @Stream.register_api()
 class sink(Sink):
@@ -67,10 +71,6 @@ def update(self, x, who=None, metadata=None):
         else:
             return []
 
-    def destroy(self):
-        super().destroy()
-        _global_sinks.remove(self)
-
 
 @Stream.register_api()
 class sink_to_textfile(Sink):

streamz/sources.py

@@ -731,3 +731,39 @@ def get_message_batch_cudf(kafka_params, topic, partition, keys, low, high, time
     finally:
         consumer.close()
     return gdf
+
+
+@Stream.register_api(staticmethod)
+class from_iterable(Source):
+    """ Emits items from an iterable.
+
+    Parameters
+    ----------
+    iterable: iterable
+        An iterable to emit messages from.
+
+    Examples
+    --------
+
+    >>> source = Stream.from_iterable(range(3))
+    >>> L = source.sink_to_list()
+    >>> source.start()
+    >>> L
+    [0, 1, 2]
+    """
+
+    def __init__(self, iterable, **kwargs):
+        super().__init__(ensure_io_loop=True, **kwargs)
+        self._iterable = iterable
+
+    def start(self):
+        self.stopped = False
+        self.loop.add_callback(self._run)
+
+    @gen.coroutine
+    def _run(self):
+        for x in self._iterable:
+            if self.stopped:
+                break
+            yield self._emit(x)
+        self.stopped = True

streamz/tests/test_sinks.py

@@ -2,7 +2,7 @@
 
 import pytest
 from streamz import Stream
-from streamz.sinks import _global_sinks
+from streamz.sinks import _global_sinks, Sink
 from streamz.utils_test import tmpfile
 
 
@@ -62,9 +62,12 @@ def test_sink_to_textfile_closes():
 
 def test_sink_destroy():
     source = Stream()
-    sink = source.sink(lambda x: None)
+    sink = Sink(source)
     ref = weakref.ref(sink)
     sink.destroy()
+
+    assert sink not in _global_sinks
+
     del sink
 
     assert ref() is None

streamz/tests/test_sources.py

@@ -95,3 +95,35 @@ def test_process():
     s.start()
     yield await_for(lambda: out == [b'0\n', b'1\n', b'2\n', b'3\n'], timeout=5)
     s.stop()
+
+
+def test_from_iterable():
+    source = Source.from_iterable(range(3))
+    L = source.sink_to_list()
+    source.start()
+    wait_for(lambda: L == [0, 1, 2], 0.1)
+
+
+def test_from_iterable_backpressure():
+    it = iter(range(5))
+    source = Source.from_iterable(it)
+    L = source.rate_limit(0.01).sink_to_list()
+    source.start()
+
+    wait_for(lambda: L == [0], 1)
+    assert next(it) == 2  # 1 is in blocked _emit
+
+
+def test_from_iterable_stop():
+    from _pytest.outcomes import Failed
+
+    source = Source.from_iterable(range(5))
+    L = source.rate_limit(0.01).sink_to_list()
+    source.start()
+
+    wait_for(lambda: L == [0], 1)
+    source.stop()
+
+    assert source.stopped
+    with pytest.raises(Failed):
+        wait_for(lambda: L == [0, 1, 2], 0.1)