Merge pull request #437 from martindurant/ws

websocket integration

Author: martindurant

ci/environment-py38.yml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - python=3.7
+  - python=3.8
   - pytest
   - flake8
   - black
@@ -32,3 +32,5 @@ dependencies:
   - flaky
   - pytest-cov
   - coveralls
+  - paho-mqtt
+  - websockets

docs/source/api.rst

@@ -43,7 +43,11 @@ Stream
 .. automethod:: Stream.emit
 .. automethod:: Stream.frequencies
 .. automethod:: Stream.register_api
+.. automethod:: Stream.sink
 .. automethod:: Stream.sink_to_list
+.. automethod:: Stream.sink_to_textfile
+.. automethod:: Stream.to_websocket
+.. automethod:: Stream.to_mqtt
 .. automethod:: Stream.update
 .. automethod:: Stream.visualize
 
@@ -55,7 +59,9 @@ Sources
    filenames
    from_kafka
    from_kafka_batched
+   from_mqtt
    from_process
+   from_websocket
    from_textfile
    from_tcp
    from_http_server

streamz/core.py

@@ -1,3 +1,4 @@
+import asyncio
 from collections import deque, defaultdict
 from datetime import timedelta
 import functools
@@ -485,15 +486,14 @@ def emit(self, x, asynchronous=False, metadata=None):
             finally:
                 thread_state.asynchronous = ts_async
         else:
-            @gen.coroutine
-            def _():
+            async def _():
                 thread_state.asynchronous = True
                 try:
-                    result = yield self._emit(x, metadata=metadata)
+                    result = await asyncio.gather(*self._emit(x, metadata=metadata))
                 finally:
                     del thread_state.asynchronous
+                return result
 
-                raise gen.Return(result)
             sync(self.loop, _)
 
     def update(self, x, who=None, metadata=None):

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,81 @@ 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_websocket(Sink):
+    """Write bytes data to websocket
+
+    The websocket will be opened on first call, and kept open. Should
+    it close at some point, future writes will fail.
+
+    Requires the ``websockets`` package.
+
+    :param uri: str
+        Something like "ws://host:port". Use "wss:" to allow TLS.
+    :param ws_kwargs: dict
+        Further kwargs to pass to ``websockets.connect``, please
+        read its documentation.
+    :param kwargs:
+        Passed to superclass
+    """
+
+    def __init__(self, upstream, uri, ws_kwargs=None, **kwargs):
+        self.uri = uri
+        self.ws_kw = ws_kwargs
+        self.ws = None
+        super().__init__(upstream, ensure_io_loop=True, **kwargs)
+
+    async def update(self, x, who=None, metadata=None):
+        import websockets
+        if self.ws is None:
+            self.ws = await websockets.connect(self.uri, **(self.ws_kw or {}))
+        await self.ws.send(x)
+
+    def destroy(self):
+        super().destroy()
+        if self.ws is not None:
+            sync(self.loop, self.ws.protocol.close)
+        self.ws = None
+
+
+@Stream.register_api()
+class to_mqtt(Sink):
+    """
+    Send data to MQTT broker
+
+    See also ``sources.from_mqtt``.
+
+    Requires ``paho.mqtt``
+
+    :param host: str
+    :param port: int
+    :param topic: str
+    :param keepalive: int
+        See mqtt docs - to keep the channel alive
+    :param client_kwargs:
+        Passed to the client's ``connect()`` method
+    """
+    def __init__(self, upstream, host, port, topic, keepalive=60, client_kwargs=None,
+                 **kwargs):
+        self.host = host
+        self.port = port
+        self.c_kw = client_kwargs or {}
+        self.client = None
+        self.topic = topic
+        self.keepalive = keepalive
+        super().__init__(upstream, ensure_io_loop=True, **kwargs)
+
+    def update(self, x, who=None, metadata=None):
+        import paho.mqtt.client as mqtt
+        if self.client is None:
+            self.client = mqtt.Client()
+            self.client.connect(self.host, self.port, self.keepalive, **self.c_kw)
+        # TODO: wait on successful delivery
+        self.client.publish(self.topic, x)
+
+    def destroy(self):
+        self.client.disconnect()
+        self.client = None
+        super().destroy()

streamz/sources.py

@@ -1,11 +1,12 @@
 import asyncio
 from glob import glob
+import queue
 import os
 import time
 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 +780,135 @@ async def run(self):
                 break
             await asyncio.gather(*self._emit(x))
         self.stopped = True
+
+
+@Stream.register_api()
+class from_websocket(Source):
+    """Read binary data from a websocket
+
+    This source will accept connections on a given port and handle messages
+    coming in.
+
+    The websockets library must be installed.
+
+    :param host: str
+        Typically "localhost"
+    :param port: int
+        Which port to listen on (must be available)
+    :param serve_kwargs: dict
+        Passed to ``websockets.serve``
+    :param kwargs:
+        Passed to superclass
+    """
+
+    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)
+
+
+@Stream.register_api()
+class from_q(Source):
+    """Source events from a threading.Queue, running another event framework
+
+    The queue is polled, i.e., there is a latency/overhead tradeoff, since
+    we cannot use ``await`` directly with a multithreaded queue.
+
+    Allows mixing of another event loop, for example pyqt, on another thread.
+    Note that, by default, a streamz.Source such as this one will start
+    an event loop in a new thread, unless otherwise specified.
+    """
+
+    def __init__(self, q, sleep_time=0.01, **kwargs):
+        """
+        :param q: threading.Queue
+            Any items pushed into here will become streamz events
+        :param sleep_time: int
+            Sets how long we wait before checking the input queue when
+            empty (in s)
+        :param kwargs:
+            passed to streamz.Source
+        """
+        self.q = q
+        self.sleep = sleep_time
+        super().__init__(**kwargs)
+
+    async def _run(self):
+        """Poll threading queue for events
+        This uses check-and-wait, but overhead is low. Could maybe have
+        a sleep-free version with an threading.Event.
+        """
+        try:
+            out = self.q.get_nowait()
+            await self.emit(out, asynchronous=True)
+        except queue.Empty:
+            await asyncio.sleep(self.sleep)
+
+
+@Stream.register_api()
+class from_mqtt(from_q):
+    """Read from MQTT source
+
+    See https://en.wikipedia.org/wiki/MQTT for a description of the protocol
+    and its uses.
+
+    See also ``sinks.to_mqtt``.
+
+    Requires ``paho.mqtt``
+
+    The outputs are ``paho.mqtt.client.MQTTMessage`` instances, which each have
+    attributes timestamp, payload, topic, ...
+
+    NB: paho.mqtt.python runs on its own thread in this implementation. We may
+    wish to instead call client.loop() directly
+
+    :param host: str
+    :param port: int
+    :param topic: str
+        (May in the future support a list of topics)
+    :param keepalive: int
+        See mqtt docs - to keep the channel alive
+    :param client_kwargs:
+        Passed to the client's ``connect()`` method
+    """
+    def __init__(self, host, port, topic, keepalive=60 , client_kwargs=None, **kwargs):
+        self.host = host
+        self.port = port
+        self.keepalive = keepalive
+        self.topic = topic
+        self.client_kwargs = client_kwargs
+        super().__init__(q=queue.Queue(), **kwargs)
+
+    def _on_connect(self, client, userdata, flags, rc):
+        client.subscribe(self.topic)
+
+    def _on_message(self, client, userdata, msg):
+        self.q.put(msg)
+
+    async def run(self):
+        import paho.mqtt.client as mqtt
+        client = mqtt.Client()
+        client.on_connect = self._on_connect
+        client.on_message = self._on_message
+        client.connect(self.host, self.port, self.keepalive, **(self.client_kwargs or {}))
+        client.loop_start()
+        await super().run()
+        client.disconnect()

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,35 @@ 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.to_websocket("ws://localhost:8989")
+    s.start()
+
+    wait_for(lambda: data == l, timeout=1)
+    s.stop()
+    s0.stop()
+
+
+def test_mqtt_roundtrip():
+    pytest.importorskip("paho.mqtt.client")
+    s0 = Stream.from_mqtt("mqtt.eclipseprojects.io", 1883, "streamz/sensor/temperature")
+    l = s0.map(lambda msg: msg.payload).sink_to_list()
+    s0.start()
+
+    data = [b'0123'] * 4
+    s = Stream.from_iterable(data)
+    s.to_mqtt("mqtt.eclipseprojects.io", 1883, "streamz/sensor/temperature")
+    s.start()
+
+    wait_for(lambda: data == l, timeout=1)
+    s.stop()
+    s0.stop()
+