Fixes #439

Author: Martin Durant

streamz/core.py

@@ -1902,89 +1902,6 @@ def cb(self):
             yield self._emit(x, self.next_metadata)
 
 
-@Stream.register_api()
-class to_kafka(Stream):
-    """ Writes data in the stream to Kafka
-
-    This stream accepts a string or bytes object. Call ``flush`` to ensure all
-    messages are pushed. Responses from Kafka are pushed downstream.
-
-    Parameters
-    ----------
-    topic : string
-        The topic which to write
-    producer_config : dict
-        Settings to set up the stream, see
-        https://docs.confluent.io/current/clients/confluent-kafka-python/#configuration
-        https://github.com/edenhill/librdkafka/blob/master/CONFIGURATION.md
-        Examples:
-        bootstrap.servers: Connection string (host:port) to Kafka
-
-    Examples
-    --------
-    >>> from streamz import Stream
-    >>> ARGS = {'bootstrap.servers': 'localhost:9092'}
-    >>> source = Stream()
-    >>> kafka = source.map(lambda x: str(x)).to_kafka('test', ARGS)
-    <to_kafka>
-    >>> for i in range(10):
-    ...     source.emit(i)
-    >>> kafka.flush()
-    """
-    def __init__(self, upstream, topic, producer_config, **kwargs):
-        import confluent_kafka as ck
-
-        self.topic = topic
-        self.producer = ck.Producer(producer_config)
-
-        kwargs["ensure_io_loop"] = True
-        Stream.__init__(self, upstream, **kwargs)
-        self.stopped = False
-        self.polltime = 0.2
-        self.loop.add_callback(self.poll)
-        self.futures = []
-
-    @gen.coroutine
-    def poll(self):
-        while not self.stopped:
-            # executes callbacks for any delivered data, in this thread
-            # if no messages were sent, nothing happens
-            self.producer.poll(0)
-            yield gen.sleep(self.polltime)
-
-    def update(self, x, who=None, metadata=None):
-        future = gen.Future()
-        self.futures.append(future)
-
-        @gen.coroutine
-        def _():
-            while True:
-                try:
-                    # this runs asynchronously, in C-K's thread
-                    self.producer.produce(self.topic, x, callback=self.cb)
-                    return
-                except BufferError:
-                    yield gen.sleep(self.polltime)
-                except Exception as e:
-                    future.set_exception(e)
-                    return
-
-        self.loop.add_callback(_)
-        return future
-
-    @gen.coroutine
-    def cb(self, err, msg):
-        future = self.futures.pop(0)
-        if msg is not None and msg.value() is not None:
-            future.set_result(None)
-            yield self._emit(msg.value())
-        else:
-            future.set_exception(err or msg.error())
-
-    def flush(self, timeout=-1):
-        self.producer.flush(timeout)
-
-
 def sync(loop, func, *args, **kwargs):
     """
     Run coroutine in loop running in separate thread.

streamz/sinks.py

@@ -109,3 +109,86 @@ def __init__(self, upstream, file, end="\n", mode="a", **kwargs):
 
     def update(self, x, who=None, metadata=None):
         self._fp.write(x + self._end)
+
+
+@Stream.register_api()
+class to_kafka(Stream):
+    """ Writes data in the stream to Kafka
+
+    This stream accepts a string or bytes object. Call ``flush`` to ensure all
+    messages are pushed. Responses from Kafka are pushed downstream.
+
+    Parameters
+    ----------
+    topic : string
+        The topic which to write
+    producer_config : dict
+        Settings to set up the stream, see
+        https://docs.confluent.io/current/clients/confluent-kafka-python/#configuration
+        https://github.com/edenhill/librdkafka/blob/master/CONFIGURATION.md
+        Examples:
+        bootstrap.servers: Connection string (host:port) to Kafka
+
+    Examples
+    --------
+    >>> from streamz import Stream
+    >>> ARGS = {'bootstrap.servers': 'localhost:9092'}
+    >>> source = Stream()
+    >>> kafka = source.map(lambda x: str(x)).to_kafka('test', ARGS)
+    <to_kafka>
+    >>> for i in range(10):
+    ...     source.emit(i)
+    >>> kafka.flush()
+    """
+    def __init__(self, upstream, topic, producer_config, **kwargs):
+        import confluent_kafka as ck
+
+        self.topic = topic
+        self.producer = ck.Producer(producer_config)
+
+        kwargs["ensure_io_loop"] = True
+        Stream.__init__(self, upstream, **kwargs)
+        self.stopped = False
+        self.polltime = 0.2
+        self.loop.add_callback(self.poll)
+        self.futures = []
+
+    @gen.coroutine
+    def poll(self):
+        while not self.stopped:
+            # executes callbacks for any delivered data, in this thread
+            # if no messages were sent, nothing happens
+            self.producer.poll(0)
+            yield gen.sleep(self.polltime)
+
+    def update(self, x, who=None, metadata=None):
+        future = gen.Future()
+        self.futures.append(future)
+
+        @gen.coroutine
+        def _():
+            while True:
+                try:
+                    # this runs asynchronously, in C-K's thread
+                    self.producer.produce(self.topic, x, callback=self.cb)
+                    return
+                except BufferError:
+                    yield gen.sleep(self.polltime)
+                except Exception as e:
+                    future.set_exception(e)
+                    return
+
+        self.loop.add_callback(_)
+        return future
+
+    @gen.coroutine
+    def cb(self, err, msg):
+        future = self.futures.pop(0)
+        if msg is not None and msg.value() is not None:
+            future.set_result(None)
+            yield self._emit(msg.value())
+        else:
+            future.set_exception(err or msg.error())
+
+    def flush(self, timeout=-1):
+        self.producer.flush(timeout)

streamz/tests/test_dask.py

@@ -1,3 +1,4 @@
+import asyncio
 from operator import add
 import random
 import time
@@ -16,21 +17,21 @@
 
 
 @gen_cluster(client=True)
-def test_map(c, s, a, b):
+async def test_map(c, s, a, b):
     source = Stream(asynchronous=True)
     futures = scatter(source).map(inc)
     futures_L = futures.sink_to_list()
     L = futures.gather().sink_to_list()
 
     for i in range(5):
-        yield source.emit(i)
+        await source.emit(i)
 
     assert L == [1, 2, 3, 4, 5]
     assert all(isinstance(f, Future) for f in futures_L)
 
 
 @gen_cluster(client=True)
-def test_map_on_dict(c, s, a, b):
+async def test_map_on_dict(c, s, a, b):
     # dask treats dicts differently, so we have to make sure
     # the user sees no difference in the streamz api.
     # Regression test against #336
@@ -43,7 +44,7 @@ def add_to_dict(d):
     L = futures.gather().sink_to_list()
 
     for i in range(5):
-        yield source.emit({"i": i})
+        await source.emit({"i": i})
 
     assert len(L) == 5
     for i, item in enumerate(sorted(L, key=lambda x: x["x"])):
@@ -52,7 +53,7 @@ def add_to_dict(d):
 
 
 @gen_cluster(client=True)
-def test_partition_then_scatter_async(c, s, a, b):
+async def test_partition_then_scatter_async(c, s, a, b):
     # Ensure partition w/ timeout before scatter works correctly for
     # asynchronous
     start = time.monotonic()
@@ -63,10 +64,10 @@ def test_partition_then_scatter_async(c, s, a, b):
 
     rc = RefCounter(loop=source.loop)
     for i in range(3):
-        yield source.emit(i, metadata=[{'ref': rc}])
+        await source.emit(i, metadata=[{'ref': rc}])
 
     while rc.count != 0 and time.monotonic() - start < 1.:
-        yield gen.sleep(1e-2)
+        await gen.sleep(1e-2)
 
     assert L == [1, 2, 3]
 
@@ -92,7 +93,7 @@ def test_partition_then_scatter_sync(loop):
 
 
 @gen_cluster(client=True)
-def test_non_unique_emit(c, s, a, b):
+async def test_non_unique_emit(c, s, a, b):
     """Regression for https://github.com/python-streamz/streams/issues/397
 
     Non-unique stream entries still need to each be processed.
@@ -103,28 +104,28 @@ def test_non_unique_emit(c, s, a, b):
 
     for _ in range(3):
         # Emit non-unique values
-        yield source.emit(0)
+        await source.emit(0)
 
     assert len(L) == 3
     assert L[0] != L[1] or L[0] != L[2]
 
 
 @gen_cluster(client=True)
-def test_scan(c, s, a, b):
+async def test_scan(c, s, a, b):
     source = Stream(asynchronous=True)
     futures = scatter(source).map(inc).scan(add)
     futures_L = futures.sink_to_list()
     L = futures.gather().sink_to_list()
 
     for i in range(5):
-        yield source.emit(i)
+        await source.emit(i)
 
     assert L == [1, 3, 6, 10, 15]
     assert all(isinstance(f, Future) for f in futures_L)
 
 
 @gen_cluster(client=True)
-def test_scan_state(c, s, a, b):
+async def test_scan_state(c, s, a, b):
     source = Stream(asynchronous=True)
 
     def f(acc, i):
@@ -133,33 +134,33 @@ def f(acc, i):
 
     L = scatter(source).scan(f, returns_state=True).gather().sink_to_list()
     for i in range(3):
-        yield source.emit(i)
+        await source.emit(i)
 
     assert L == [0, 1, 3]
 
 
 @gen_cluster(client=True)
-def test_zip(c, s, a, b):
+async def test_zip(c, s, a, b):
     a = Stream(asynchronous=True)
     b = Stream(asynchronous=True)
     c = scatter(a).zip(scatter(b))
 
     L = c.gather().sink_to_list()
 
-    yield a.emit(1)
-    yield b.emit('a')
-    yield a.emit(2)
-    yield b.emit('b')
+    await a.emit(1)
+    await b.emit('a')
+    await a.emit(2)
+    await b.emit('b')
 
     assert L == [(1, 'a'), (2, 'b')]
 
 
 @gen_cluster(client=True)
-def test_accumulate(c, s, a, b):
+async def test_accumulate(c, s, a, b):
     source = Stream(asynchronous=True)
     L = source.scatter().accumulate(lambda acc, x: acc + x, with_state=True).gather().sink_to_list()
     for i in range(3):
-        yield source.emit(i)
+        await source.emit(i)
     assert L[-1][1] == 3
 
 
@@ -169,10 +170,9 @@ def test_sync(loop):  # noqa: F811
             source = Stream()
             L = source.scatter().map(inc).gather().sink_to_list()
 
-            @gen.coroutine
-            def f():
+            async def f():
                 for i in range(10):
-                    yield source.emit(i, asynchronous=True)
+                    await source.emit(i, asynchronous=True)
 
             sync(loop, f)
 
@@ -193,24 +193,24 @@ def test_sync_2(loop):  # noqa: F811
 
 
 @gen_cluster(client=True, nthreads=[('127.0.0.1', 1)] * 2)
-def test_buffer(c, s, a, b):
+async def test_buffer(c, s, a, b):
     source = Stream(asynchronous=True)
     L = source.scatter().map(slowinc, delay=0.5).buffer(5).gather().sink_to_list()
 
     start = time.time()
     for i in range(5):
-        yield source.emit(i)
+        await source.emit(i)
     end = time.time()
     assert end - start < 0.5
 
     for i in range(5, 10):
-        yield source.emit(i)
+        await source.emit(i)
 
     end2 = time.time()
     assert end2 - start > (0.5 / 3)
 
     while len(L) < 10:
-        yield gen.sleep(0.01)
+        await gen.sleep(0.01)
         assert time.time() - start < 5
 
     assert L == list(map(inc, range(10)))
@@ -242,7 +242,7 @@ def test_buffer_sync(loop):  # noqa: F811
 
 
 @pytest.mark.xfail(reason='')
-def test_stream_shares_client_loop(loop):  # noqa: F811
+async def test_stream_shares_client_loop(loop):  # noqa: F811
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as client:  # noqa: F841
             source = Stream()
@@ -251,14 +251,14 @@ def test_stream_shares_client_loop(loop):  # noqa: F811
 
 
 @gen_cluster(client=True)
-def test_starmap(c, s, a, b):
+async def test_starmap(c, s, a, b):
     def add(x, y, z=0):
         return x + y + z
 
     source = Stream(asynchronous=True)
     L = source.scatter().starmap(add, z=10).gather().sink_to_list()
 
     for i in range(5):
-        yield source.emit((i, i))
+        await source.emit((i, i))
 
     assert L == [10, 12, 14, 16, 18]