refactor and make less tornado

Author: Martin Durant

streamz/compatibility.py

@@ -23,7 +23,7 @@
 
 from collections.abc import Iterable
 
-from .compatibility import get_thread_identity
+from threading import get_ident as get_thread_identity
 from .orderedweakset import OrderedWeakrefSet
 
 no_default = '--no-default--'

streamz/core.py

@@ -1,6 +1,7 @@
+import asyncio
 from glob import glob
 import os
-
+import inspect
 import time
 import tornado.ioloop
 from tornado import gen
@@ -35,15 +36,26 @@ def write(text):
 class Source(Stream):
     _graphviz_shape = 'doubleoctagon'
 
-    def __init__(self, **kwargs):
+    def __init__(self, start=False, **kwargs):
         self.stopped = True
-        super(Source, self).__init__(**kwargs)
+        super().__init__(ensure_io_loop=True, **kwargs)
+        if start:
+            self.start()
 
-    def stop(self):  # pragma: no cover
-        # fallback stop method - for poll functions with while not self.stopped
+    def stop(self):
+        """set self.stopped, which will cause polling to stop after next run"""
         if not self.stopped:
             self.stopped = True
 
+    def start(self):
+        """start polling"""
+        self.stopped = False
+        self.loop.add_callback(self.run)
+
+    async def run(self):
+        while not self.stopped:
+            await self._run()
+
 
 @Stream.register_api(staticmethod)
 class from_textfile(Source):
@@ -74,44 +86,32 @@ class from_textfile(Source):
     -------
     Stream
     """
-    def __init__(self, f, poll_interval=0.100, delimiter='\n', start=False,
+    def __init__(self, f, poll_interval=0.100, delimiter='\n',
                  from_end=False, **kwargs):
         if isinstance(f, str):
             f = open(f)
+        self.buffer = ''
+        if self.from_end:
+            # this only happens when we are ready to read
+            self.file.seek(0, 2)
         self.file = f
         self.from_end = from_end
         self.delimiter = delimiter
 
         self.poll_interval = poll_interval
-        super(from_textfile, self).__init__(ensure_io_loop=True, **kwargs)
-        self.stopped = True
-        self.started = False
-        if start:
-            self.start()
+        super().__init__(**kwargs)
 
-    def start(self):
-        self.stopped = False
-        self.started = False
-        self.loop.add_callback(self.do_poll)
-
-    @gen.coroutine
-    def do_poll(self):
-        buffer = ''
-        if self.from_end:
-            # this only happens when we are ready to read
-            self.file.seek(0, 2)
-        while not self.stopped:
-            self.started = True
-            line = self.file.read()
-            if line:
-                buffer = buffer + line
-                if self.delimiter in buffer:
-                    parts = buffer.split(self.delimiter)
-                    buffer = parts.pop(-1)
-                    for part in parts:
-                        yield self._emit(part + self.delimiter)
-            else:
-                yield gen.sleep(self.poll_interval)
+    async def _run(self):
+        line = self.file.read()
+        if line:
+            self.buffer = self.buffer + line
+            if self.delimiter in self.buffer:
+                parts = self.buffer.split(self.delimiter)
+                self.buffer = parts.pop(-1)
+                for part in parts:
+                    await self._emit(part + self.delimiter)
+        else:
+            await asyncio.sleep(self.poll_interval)
 
 
 @Stream.register_api(staticmethod)
@@ -133,7 +133,7 @@ class filenames(Source):
     >>> source = Stream.filenames('path/to/dir')  # doctest: +SKIP
     >>> source = Stream.filenames('path/to/*.csv', poll_interval=0.500)  # doctest: +SKIP
     """
-    def __init__(self, path, poll_interval=0.100, start=False, **kwargs):
+    def __init__(self, path, poll_interval=0.100, **kwargs):
         if '*' not in path:
             if os.path.isdir(path):
                 if not path.endswith(os.path.sep):
@@ -142,26 +142,15 @@ def __init__(self, path, poll_interval=0.100, start=False, **kwargs):
         self.path = path
         self.seen = set()
         self.poll_interval = poll_interval
-        self.stopped = True
-        super(filenames, self).__init__(ensure_io_loop=True)
-        if start:
-            self.start()
+        super().__init__(**kwargs)
 
-    def start(self):
-        self.stopped = False
-        self.loop.add_callback(self.do_poll)
-
-    @gen.coroutine
-    def do_poll(self):
-        while True:
-            filenames = set(glob(self.path))
-            new = filenames - self.seen
-            for fn in sorted(new):
-                self.seen.add(fn)
-                yield self._emit(fn)
-            yield gen.sleep(self.poll_interval)  # TODO: remove poll if delayed
-            if self.stopped:
-                break
+    async def _run(self):
+        filenames = set(glob(self.path))
+        new = filenames - self.seen
+        for fn in sorted(new):
+            self.seen.add(fn)
+            await self._emit(fn)
+        await asyncio.sleep(self.poll_interval)  # TODO: remove poll if delayed
 
 
 @Stream.register_api(staticmethod)
@@ -191,42 +180,31 @@ class from_tcp(Source):
 
     >>> source = Source.from_tcp(4567)  # doctest: +SKIP
     """
-    def __init__(self, port, delimiter=b'\n', start=False,
-                 server_kwargs=None):
-        super(from_tcp, self).__init__(ensure_io_loop=True)
-        self.stopped = True
+    def __init__(self, port, delimiter=b'\n', server_kwargs=None, **kwargs):
         self.server_kwargs = server_kwargs or {}
         self.port = port
         self.server = None
         self.delimiter = delimiter
-        if start:  # pragma: no cover
-            self.start()
+        super().__init__(**kwargs)
 
-    @gen.coroutine
-    def _start_server(self):
+    def run(self):
         from tornado.tcpserver import TCPServer
         from tornado.iostream import StreamClosedError
 
         class EmitServer(TCPServer):
             source = self
 
-            @gen.coroutine
-            def handle_stream(self, stream, address):
-                while True:
+            async def handle_stream(self, stream, address):
+                while not self.source.stopped:
                     try:
-                        data = yield stream.read_until(self.source.delimiter)
-                        yield self.source._emit(data)
+                        data = await stream.read_until(self.source.delimiter)
+                        await self.source._emit(data)
                     except StreamClosedError:
                         break
 
         self.server = EmitServer(**self.server_kwargs)
         self.server.listen(self.port)
 
-    def start(self):
-        if self.stopped:
-            self.loop.add_callback(self._start_server)
-            self.stopped = False
-
     def stop(self):
         if not self.stopped:
             self.server.stop()
@@ -260,26 +238,22 @@ class from_http_server(Source):
 
     """
 
-    def __init__(self, port, path='/.*', start=False, server_kwargs=None):
+    def __init__(self, port, path='/.*', server_kwargs=None, **kwargs):
         self.port = port
         self.path = path
         self.server_kwargs = server_kwargs or {}
-        super(from_http_server, self).__init__(ensure_io_loop=True)
-        self.stopped = True
         self.server = None
-        if start:  # pragma: no cover
-            self.start()
+        super().__init__(**kwargs)
 
-    def _start_server(self):
+    def run(self):
         from tornado.web import Application, RequestHandler
         from tornado.httpserver import HTTPServer
 
         class Handler(RequestHandler):
             source = self
 
-            @gen.coroutine
-            def post(self):
-                yield self.source._emit(self.request.body)
+            async def post(self):
+                await asyncio.gather(*self.source._emit(self.request.body))
                 self.write('OK')
 
         application = Application([
@@ -288,12 +262,6 @@ def post(self):
         self.server = HTTPServer(application, **self.server_kwargs)
         self.server.listen(self.port)
 
-    def start(self):
-        """Start HTTP server and listen"""
-        if self.stopped:
-            self.loop.add_callback(self._start_server)
-            self.stopped = False
-
     def stop(self):
         """Shutdown HTTP server"""
         if not self.stopped:
@@ -325,46 +293,36 @@ class from_process(Source):
     >>> source = Source.from_process(['ping', 'localhost'])  # doctest: +SKIP
     """
 
-    def __init__(self, cmd, open_kwargs=None, with_stderr=False, start=False):
+    def __init__(self, cmd, open_kwargs=None, with_stderr=False, with_end=True,
+                 **kwargs):
         self.cmd = cmd
         self.open_kwargs = open_kwargs or {}
         self.with_stderr = with_stderr
-        super(from_process, self).__init__(ensure_io_loop=True)
-        self.stopped = True
+        self.with_end = with_end
         self.process = None
-        if start:  # pragma: no cover
-            self.start()
+        super().__init__(**kwargs)
 
-    @gen.coroutine
-    def _start_process(self):
-        # should be done in asyncio (py3 only)? Apparently can handle Windows
-        # with appropriate config.
-        from tornado.process import Subprocess
-        from tornado.iostream import StreamClosedError
+    async def run(self):
         import subprocess
-        stderr = subprocess.STDOUT if self.with_stderr else subprocess.PIPE
-        process = Subprocess(self.cmd, stdout=Subprocess.STREAM,
-                             stderr=stderr, **self.open_kwargs)
+        stderr = subprocess.STDOUT if self.with_stderr else None
+        if isinstance(self.cmd, (list, tuple)):
+            cmd, *args = self.cmd
+        else:
+            cmd, args = self.cmd, ()
+        process = await asyncio.create_subprocess_exec(
+            cmd, *args, stdout=subprocess.PIPE,
+            stderr=stderr, **self.open_kwargs)
         while not self.stopped:
             try:
-                out = yield process.stdout.read_until(b'\n')
-            except StreamClosedError:
-                # process exited
-                break
-            yield self._emit(out)
-        yield process.stdout.close()
-        process.proc.terminate()
-
-    def start(self):
-        """Start external process"""
-        if self.stopped:
-            self.loop.add_callback(self._start_process)
-            self.stopped = False
-
-    def stop(self):
-        """Shutdown external process"""
-        if not self.stopped:
-            self.stopped = True
+                out = await process.stdout.readuntil(b'\n')
+            except asyncio.IncompleteReadError as err:
+                if self.with_end:
+                    out = err.partial
+                else:
+                    break
+            await asyncio.gather(*self._emit(out))
+        process.terminate()
+        await process.wait()
 
 
 @Stream.register_api(staticmethod)
@@ -401,15 +359,12 @@ class from_kafka(Source):
     ...            'group.id': 'streamz'})  # doctest: +SKIP
 
     """
-    def __init__(self, topics, consumer_params, poll_interval=0.1, start=False, **kwargs):
+    def __init__(self, topics, consumer_params, poll_interval=0.1, **kwargs):
         self.cpars = consumer_params
         self.consumer = None
         self.topics = topics
         self.poll_interval = poll_interval
-        super(from_kafka, self).__init__(ensure_io_loop=True, **kwargs)
-        self.stopped = True
-        if start:
-            self.start()
+        super().__init__(**kwargs)
 
     def do_poll(self):
         if self.consumer is not None:
@@ -470,10 +425,9 @@ def __init__(self, topic, consumer_params, poll_interval='1s',
         self.max_batch_size = max_batch_size
         self.keys = keys
         self.engine = engine
-        self.stopped = True
         self.started = False
 
-        super(FromKafkaBatched, self).__init__(ensure_io_loop=True, **kwargs)
+        super().__init__(**kwargs)
 
     @gen.coroutine
     def poll_kafka(self):
@@ -753,17 +707,12 @@ class from_iterable(Source):
     """
 
     def __init__(self, iterable, **kwargs):
-        super().__init__(ensure_io_loop=True, **kwargs)
         self._iterable = iterable
+        super().__init__(**kwargs)
 
-    def start(self):
-        self.stopped = False
-        self.loop.add_callback(self._run)
-
-    @gen.coroutine
-    def _run(self):
+    async def run(self):
         for x in self._iterable:
             if self.stopped:
                 break
-            yield self._emit(x)
+            await asyncio.gather(*self._emit(x))
         self.stopped = True

streamz/sources.py

@@ -86,11 +86,19 @@ def test_http():
         requests.post('http://localhost:%i/other' % port, data=b'data2')
 
 
-@flaky(max_runs=3, min_passes=1)
+#@flaky(max_runs=3, min_passes=1)
 @gen_test(timeout=60)
 def test_process():
+    import sys
+    import asyncio
     cmd = ["python", "-c", "for i in range(4): print(i)"]
     s = Source.from_process(cmd)
+    if sys.platform != "win32":
+        # don't know why - something with pytest and new processes
+        policy = asyncio.get_event_loop_policy()
+        watcher = asyncio.SafeChildWatcher()
+        policy.set_child_watcher(watcher)
+        watcher.attach_loop(s.loop.asyncio_loop)
     out = s.sink_to_list()
     s.start()
     yield await_for(lambda: out == [b'0\n', b'1\n', b'2\n', b'3\n'], timeout=5)