imporved docs, comments, updated tests

Author: roveo

streamz/sinks.py

@@ -1,7 +1,11 @@
+import inspect
+import weakref
+
 from tornado import gen
 
 from streamz import Stream
 
+# sinks add themselves here to avoid being garbage-collected
 _global_sinks = set()
 
 
@@ -18,6 +22,16 @@ def __init__(self, upstream, **kwargs):
 class sink(Sink):
     """ Apply a function on every element
 
+    Parameters
+    ----------
+    func: callable
+        A function that will be applied on every element.
+    args:
+        Positional arguments that will be passed to ``func`` after the incoming element.
+    kwargs:
+        Stream-specific arguments will be passed to ``Stream.__init__``, the rest of
+        them will be passed to ``func``.
+
     Examples
     --------
     >>> source = Stream()
@@ -40,10 +54,11 @@ class sink(Sink):
     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
+        sig = set(inspect.signature(Stream).parameters)
+        stream_kwargs = {k: v for (k, v) in kwargs.items() if k in sig}
+        self.kwargs = {k: v for (k, v) in kwargs.items() if k not in sig}
         self.args = args
-        super().__init__(upstream, stream_name=stream_name)
+        super().__init__(upstream, **stream_kwargs)
 
     def update(self, x, who=None, metadata=None):
         result = self.func(x, *self.args, **self.kwargs)
@@ -59,8 +74,8 @@ class sink_to_textfile(Sink):
 
         Type of elements must be ``str``.
 
-        Arguments
-        ---------
+        Parameters
+        ----------
         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
@@ -83,14 +98,14 @@ class sink_to_textfile(Sink):
         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
+        self._fp = open(file, mode=mode) if isinstance(file, str) else file
+        weakref.finalize(self, self._fp.close)
         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)
+        yield self.loop.run_in_executor(None, self._fp.write, x + self._end)

streamz/tests/test_sinks.py

@@ -1,5 +1,3 @@
-from time import sleep
-
 import pytest
 from streamz import Stream
 from streamz.sinks import _global_sinks
@@ -16,33 +14,34 @@ def mycustomsink(elem, key, prefix=""):
         L[key].append(elem)
 
     s = Stream()
-    s.sink(mycustomsink, "cat", "super")
-
+    sink = s.sink(mycustomsink, "cat", "super", stream_name="test")
     s.emit(1)
     s.emit(2)
+
     assert L['supercat'] == [1, 2]
+    assert sink.name == "test"
 
 
 def test_sink_to_textfile_fp():
     source = Stream()
-    with tmpfile() as filename, open(filename, "w", buffering=1) as fp:
+    with tmpfile() as filename, open(filename, "w") as fp:
         source.map(str).sink_to_textfile(fp)
         source.emit(0)
         source.emit(1)
 
-        sleep(0.01)
+        fp.flush()
 
         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)
+        _sink = source.map(str).sink_to_textfile(filename)
         source.emit(0)
         source.emit(1)
 
-        sleep(0.01)
+        _sink._fp.flush()
 
         assert open(filename, "r").read() == "0\n1\n"
 
@@ -55,5 +54,5 @@ def test_sink_to_textfile_closes():
         _global_sinks.remove(sink)
         del sink
 
-        with pytest.raises(ValueError):  # I/O operation on closed file
+        with pytest.raises(ValueError, match=r"I/O operation on closed file\."):
             fp.write(".")