refactor sinks, sink_to_textfile

Author: roveo

.travis.yml

@@ -21,7 +21,7 @@ install:
   - conda update conda
 
   # Install dependencies
-  - conda env create --name test-streamz --file ./conda/environments/streamz_dev.yml
+  - travis_wait 30 conda env create --name test-streamz --file ./conda/environments/streamz_dev.yml
   - source activate test-streamz
 
   - python setup.py install

streamz/__init__.py

@@ -3,6 +3,7 @@
 from .core import *
 from .graph import *
 from .sources import *
+from .sinks import *
 try:
     from .dask import DaskStream, scatter
 except ImportError:

streamz/core.py

@@ -27,8 +27,6 @@
 
 no_default = '--no-default--'
 
-_global_sinks = set()
-
 _html_update_streams = set()
 
 thread_state = threading.local()
@@ -622,48 +620,6 @@ def _release_refs(self, metadata, n=1):
                 m['ref'].release(n)
 
 
-@Stream.register_api()
-class sink(Stream):
-    """ Apply a function on every element
-
-    Examples
-    --------
-    >>> source = Stream()
-    >>> L = list()
-    >>> source.sink(L.append)
-    >>> source.sink(print)
-    >>> source.sink(print)
-    >>> source.emit(123)
-    123
-    123
-    >>> L
-    [123]
-
-    See Also
-    --------
-    map
-    Stream.sink_to_list
-    """
-    _graphviz_shape = 'trapezium'
-
-    def __init__(self, upstream, func, *args, **kwargs):
-        self.func = func
-        # take the stream specific kwargs out
-        stream_name = kwargs.pop("stream_name", None)
-        self.kwargs = kwargs
-        self.args = args
-
-        Stream.__init__(self, upstream, stream_name=stream_name)
-        _global_sinks.add(self)
-
-    def update(self, x, who=None, metadata=None):
-        result = self.func(x, *self.args, **self.kwargs)
-        if gen.isawaitable(result):
-            return result
-        else:
-            return []
-
-
 @Stream.register_api()
 class map(Stream):
     """ Apply a function to every element in the stream

streamz/sinks.py

@@ -0,0 +1,95 @@
+from tornado import gen
+
+from streamz import Stream
+
+_global_sinks = set()
+
+
+class Sink(Stream):
+
+    _graphviz_shape = 'trapezium'
+
+    def __init__(self, upstream, **kwargs):
+        super().__init__(upstream, **kwargs)
+        _global_sinks.add(self)
+
+
+@Stream.register_api()
+class sink(Sink):
+    """ Apply a function on every element
+
+    Examples
+    --------
+    >>> source = Stream()
+    >>> L = list()
+    >>> source.sink(L.append)
+    >>> source.sink(print)
+    >>> source.sink(print)
+    >>> source.emit(123)
+    123
+    123
+    >>> L
+    [123]
+
+    See Also
+    --------
+    map
+    Stream.sink_to_list
+    """
+
+    def __init__(self, upstream, func, *args, **kwargs):
+        self.func = func
+        # take the stream specific kwargs out
+        stream_name = kwargs.pop("stream_name", None)
+        self.kwargs = kwargs
+        self.args = args
+        super().__init__(upstream, stream_name=stream_name)
+
+    def update(self, x, who=None, metadata=None):
+        result = self.func(x, *self.args, **self.kwargs)
+        if gen.isawaitable(result):
+            return result
+        else:
+            return []
+
+
+@Stream.register_api()
+class sink_to_textfile(Sink):
+    """ Write elements to a plain text file, one element per line. Type of elements
+        must be ``str``.
+
+        Arguments
+        ---------
+        file: str or file-like
+            File to write the elements to. ``str`` is treated as a file name to open.
+            If file-like, descriptor must be open in text mode.
+            Note that the file descriptor will be closed when sink is destroyed.
+        end: str, optional
+            This value will be written to the file after each element.
+            Defaults to newline character.
+        mode: str, optional
+            If file is ``str``, file will be opened in this mode. Defaults to ``"a"``
+            (append mode).
+
+        Examples
+        --------
+        >>> source = Stream()
+        >>> source.map(str).sink_to_file("test.txt")
+        >>> source.emit(0)
+        >>> source.emit(1)
+        >>> print(open("test.txt", "r").read())
+        0
+        1
+    """
+    def __init__(self, upstream, file, end="\n", mode="a", **kwargs):
+        self._fp = open(file, mode=mode, buffering=1) if isinstance(file, str) else file
+        self._end = end
+        super().__init__(upstream, ensure_io_loop=True, **kwargs)
+
+    def __del__(self):
+        self._fp.close()
+
+    @gen.coroutine
+    def update(self, x, who=None, metadata=None):
+        yield self.loop.run_in_executor(None, self._fp.write, x)
+        yield self.loop.run_in_executor(None, self._fp.write, self._end)

streamz/tests/test_core.py

@@ -336,23 +336,6 @@ def test_sink_to_file():
         assert data == 'a\nb\n'
 
 
-def test_sink_with_args_and_kwargs():
-    L = dict()
-
-    def mycustomsink(elem, key, prefix=""):
-        key = prefix + key
-        if key not in L:
-            L[key] = list()
-        L[key].append(elem)
-
-    s = Stream()
-    s.sink(mycustomsink, "cat", "super")
-
-    s.emit(1)
-    s.emit(2)
-    assert L['supercat'] == [1, 2]
-
-
 @gen_test()
 def test_counter():
     counter = itertools.count()

streamz/tests/test_sinks.py

@@ -0,0 +1,45 @@
+from time import sleep
+
+from streamz import Stream
+from streamz.utils_test import tmpfile
+
+
+def test_sink_with_args_and_kwargs():
+    L = dict()
+
+    def mycustomsink(elem, key, prefix=""):
+        key = prefix + key
+        if key not in L:
+            L[key] = list()
+        L[key].append(elem)
+
+    s = Stream()
+    s.sink(mycustomsink, "cat", "super")
+
+    s.emit(1)
+    s.emit(2)
+    assert L['supercat'] == [1, 2]
+
+
+def test_sink_to_textfile_fp():
+    source = Stream()
+    with tmpfile() as filename, open(filename, "w", buffering=1) as fp:
+        source.map(str).sink_to_textfile(fp)
+        source.emit(0)
+        source.emit(1)
+
+        sleep(0.01)
+
+        assert open(filename, "r").read() == "0\n1\n"
+
+
+def test_sink_to_textfile_named():
+    source = Stream()
+    with tmpfile() as filename:
+        source.map(str).sink_to_textfile(filename)
+        source.emit(0)
+        source.emit(1)
+
+        sleep(0.01)
+
+        assert open(filename, "r").read() == "0\n1\n"