websocket integration

Author: Martin Durant

streamz/sinks.py

@@ -4,6 +4,7 @@
 from tornado import gen
 
 from streamz import Stream
+from streamz.core import sync
 
 # sinks add themselves here to avoid being garbage-collected
 _global_sinks = set()
@@ -109,3 +110,26 @@ 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 sink_to_websocket(Sink):
+
+    def __init__(self, upstream, uri, ws_kwargs=None, **kwargs):
+        self.uri = uri
+        self.ws_kw = ws_kwargs
+        self.ws = None
+        super().__init__(upstream, **kwargs)
+
+    @gen.coroutine
+    def update(self, x, who=None, metadata=None):
+        import websockets
+        if self.ws is None:
+            self.ws = yield websockets.connect(self.uri, **(self.ws_kw or {}))
+        yield self.ws.send(x)
+
+    def destroy(self):
+        super(sink_to_websocket, self).destroy()
+        if self.ws is not None:
+            sync(self.loop, self.ws.protocol.close)
+        self.ws = None

streamz/sources.py

@@ -5,7 +5,7 @@
 from tornado import gen
 import weakref
 
-from .core import Stream, convert_interval, RefCounter
+from .core import Stream, convert_interval, RefCounter, sync
 
 
 def sink_to_file(filename, upstream, mode='w', prefix='', suffix='\n', flush=False):
@@ -779,3 +779,29 @@ async def run(self):
                 break
             await asyncio.gather(*self._emit(x))
         self.stopped = True
+
+
+@Stream.register_api()
+class from_websocket(Source):
+    def __init__(self, host, port, serve_kwargs=None, **kwargs):
+        self.host = host
+        self.port = port
+        self.s_kw = serve_kwargs
+        self.server = None
+        super().__init__(**kwargs)
+
+    @gen.coroutine
+    def _read(self, ws, path):
+        while not self.stopped:
+            data = yield ws.recv()
+            yield self._emit(data)
+
+    async def run(self):
+        import websockets
+        self.server = await websockets.serve(
+            self._read, self.host, self.port, **(self.s_kw or {})
+        )
+
+    def stop(self):
+        self.server.close()
+        sync(self.loop, self.server.wait_closed)

streamz/tests/test_sinks.py

@@ -3,7 +3,7 @@
 import pytest
 from streamz import Stream
 from streamz.sinks import _global_sinks, Sink
-from streamz.utils_test import tmpfile
+from streamz.utils_test import tmpfile, wait_for
 
 
 def test_sink_with_args_and_kwargs():
@@ -71,3 +71,18 @@ def test_sink_destroy():
     del sink
 
     assert ref() is None
+
+
+def test_ws_roundtrip():
+    pytest.importorskip("websockets")
+    s0 = Stream.from_websocket("localhost", 8989, start=True)
+    l = s0.sink_to_list()
+
+    data = [b'0123'] * 4
+    s = Stream.from_iterable(data)
+    s.sink_to_websocket("ws://localhost:8989")
+    s.start()
+
+    wait_for(lambda: data == l, timeout=1)
+    s.stop()
+    s0.stop()