A simple plugin system

.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

docs/source/index.rst

@@ -122,3 +122,4 @@ data streaming systems like `Apache Flink <https://flink.apache.org/>`_,
    collections-api.rst
    async.rst
    plotting.rst
+   plugins.rst

docs/source/plugins.rst

@@ -0,0 +1,102 @@
+Plugins
+=======
+
+In addition to using ``@Stream.register_api()`` decorator, custom stream nodes can
+be added to Streamz by installing 3rd-party Python packages.
+
+
+Known plugins
+-------------
+
+Extras
+++++++
+
+These plugins are supported by the Streamz community and can be installed as extras,
+e.g. ``pip install streamz[kafka]``.
+
+There are no plugins here yet, but hopefully soon there will be.
+
+.. only:: comment
+    ================= ======================================================
+    Extra name        Description
+    ================= ======================================================
+    ``files``         Advanced filesystem operations: listening for new
+                    files in a directory, writing to multiple files etc.
+    ``kafka``         Reading from and writing to Kafka topics.
+    ================= ======================================================
+
+
+Entry points
+------------
+
+Plugins register themselves with Streamz by using ``entry_points`` argument
+in ``setup.py``:
+
+.. code-block:: Python
+
+    # setup.py
+
+    from setuptools import setup
+
+    setup(
+        name="streamz_example_plugin",
+        version="0.0.1",
+        entry_points={
+            "streamz.nodes": [
+                "repeat = streamz_example_plugin:RepeatNode"
+            ]
+        }
+    )
+
+In this example, ``RepeatNode`` class will be imported from
+``streamz_example_plugin`` package and will be available as ``Stream.repeat``.
+In practice, class name and entry point name (the part before ``=`` in entry point
+definition) are usually the same, but they `can` be different.
+
+Different kinds of add-ons go into different entry point groups:
+
+=========== ======================= =====================
+ Node type   Required parent class   Entry point group
+=========== ======================= =====================
+ Source      ``streamz.Source``      ``streamz.sources``
+ Node        ``streamz.Stream``      ``streamz.nodes``
+ Sink        ``streamz.Stream``      ``streamz.sinks``
+=========== ======================= =====================
+
+
+Lazy loading
+++++++++++++
+
+Streamz will attach methods from existing plugins to the ``Stream`` class when it's
+imported, but actual classes will be loaded only when the corresponding ``Stream``
+method is first called. Streamz will also validate the loaded class before attaching it
+and will raise an appropriate exception if validation fails.
+
+
+Reference implementation
+------------------------
+
+Let's look at how stream nodes can be implemented.
+
+.. code-block:: Python
+
+    # __init__.py
+
+    from tornado import gen
+    from streamz import Stream
+
+
+    class RepeatNode(Stream):
+
+        def __init__(self, upstream, n, **kwargs):
+            super().__init__(upstream, ensure_io_loop=True, **kwargs)
+            self._n = n
+
+        @gen.coroutine
+        def update(self, x, who=None, metadata=None):
+            for _ in range(self._n):
+                yield self._emit(x, metadata=metadata)
+
+As you can see, implementation is the same as usual, but there's no
+``@Stream.register_api()`` — Streamz will take care of that when loading the plugin.
+It will still work if you add the decorator, but you don't have to.

requirements.txt

@@ -2,3 +2,4 @@ tornado
 toolz
 zict
 six
+setuptools

streamz/__init__.py

@@ -3,6 +3,10 @@
 from .core import *
 from .graph import *
 from .sources import *
+from .plugins import load_plugins
+
+load_plugins(Stream)
+
 try:
     from .dask import DaskStream, scatter
 except ImportError:

streamz/core.py

@@ -260,7 +260,7 @@ def _remove_upstream(self, upstream):
             self.upstreams.remove(upstream)
 
     @classmethod
-    def register_api(cls, modifier=identity):
+    def register_api(cls, modifier=identity, attribute_name=None):
         """ Add callable to Stream API
 
         This allows you to register a new method onto this class.  You can use
@@ -273,7 +273,7 @@ def register_api(cls, modifier=identity):
             >>> Stream().foo(...)  # this works now
 
         It attaches the callable as a normal attribute to the class object.  In
-        doing so it respsects inheritance (all subclasses of Stream will also
+        doing so it respects inheritance (all subclasses of Stream will also
         get the foo attribute).
 
         By default callables are assumed to be instance methods.  If you like
@@ -285,15 +285,41 @@ def register_api(cls, modifier=identity):
             ...     ...
 
             >>> Stream.foo(...)  # Foo operates as a static method
+
+        You can also provide an optional ``attribute_name`` argument to control
+        the name of the attribute your callable will be attached as.
+
+            >>> @Stream.register_api(attribute_name="bar")
+            ... class foo(Stream):
+            ...     ...
+
+            >> Stream().bar(...)  # foo was actually attached as bar
         """
         def _(func):
             @functools.wraps(func)
             def wrapped(*args, **kwargs):
                 return func(*args, **kwargs)
-            setattr(cls, func.__name__, modifier(wrapped))
+            name = attribute_name if attribute_name else func.__name__
+            setattr(cls, name, modifier(wrapped))
             return func
         return _
 
+    @classmethod
+    def register_plugin_entry_point(cls, entry_point, modifier=identity):
+        def stub(*args, **kwargs):
+            node = entry_point.load()
+            if not issubclass(node, Stream):
+                raise TypeError(
+                    f"Error loading {entry_point.name} "
+                    f"from module {entry_point.module_name}: "
+                    f"{node.__class__.__name__} must be a subclass of Stream"
+                )
+            cls.register_api(
+                modifier=modifier, attribute_name=entry_point.name
+            )(node)
+            return node(*args, **kwargs)
+        cls.register_api(modifier=modifier, attribute_name=entry_point.name)(stub)
+
     def start(self):
         """ Start any upstream sources """
         for upstream in self.upstreams:

streamz/dataframe/tests/test_dataframes.py

@@ -1008,9 +1008,9 @@ def test_windowed_groupby_aggs_with_start_state(stream):
     out_df1 = pd.DataFrame({'name':['Alice', 'Bob', 'Linda', 'Tom'], 'amount':[50.0, 550.0, 100.0, 150.0]})
     assert_eq(output1[-1][1].reset_index(), out_df1)
 
-    
+
 def test_dir(stream):
     example = pd.DataFrame({'name': [], 'amount': []})
     sdf = DataFrame(stream, example=example)
     assert 'name' in dir(sdf)
-    assert 'amount' in dir(sdf)    
+    assert 'amount' in dir(sdf)

streamz/plugins.py

@@ -0,0 +1,10 @@
+import pkg_resources
+
+
+def load_plugins(cls):
+    for entry_point in pkg_resources.iter_entry_points("streamz.sources"):
+        cls.register_plugin_entry_point(entry_point, staticmethod)
+    for entry_point in pkg_resources.iter_entry_points("streamz.nodes"):
+        cls.register_plugin_entry_point(entry_point)
+    for entry_point in pkg_resources.iter_entry_points("streamz.sinks"):
+        cls.register_plugin_entry_point(entry_point)

streamz/tests/test_plugins.py

@@ -0,0 +1,69 @@
+import pytest
+from streamz import Source, Stream
+
+
+class MockEntryPoint:
+
+    def __init__(self, name, cls, module_name=None):
+        self.name = name
+        self.cls = cls
+        self.module_name = module_name
+
+    def load(self):
+        return self.cls
+
+
+def test_register_plugin_entry_point():
+    class test_stream(Stream):
+        pass
+
+    entry_point = MockEntryPoint("test_node", test_stream)
+    Stream.register_plugin_entry_point(entry_point)
+
+    assert Stream.test_node.__name__ == "stub"
+
+    Stream().test_node()
+
+    assert Stream.test_node.__name__ == "test_stream"
+
+
+def test_register_plugin_entry_point_modifier():
+    class test_source(Source):
+        pass
+
+    def modifier(fn):
+        fn.__name__ = 'modified_name'
+        return staticmethod(fn)
+
+    entry_point = MockEntryPoint("from_test", test_source)
+    Stream.register_plugin_entry_point(entry_point, modifier)
+
+    Stream.from_test()
+
+    assert Stream.from_test.__name__ == "modified_name"
+
+
+def test_register_plugin_entry_point_raises():
+    class invalid_node:
+        pass
+
+    entry_point = MockEntryPoint("test", invalid_node, "test_module.test")
+
+    Stream.register_plugin_entry_point(entry_point)
+
+    with pytest.raises(TypeError):
+        Stream.test()
+
+
+def test_register_plugin_entry_point_already_registered():
+    @Stream.register_api()
+    class test(Stream):
+        pass
+
+    entry_point = MockEntryPoint("test_double", test, "test_module")
+
+    Stream.register_plugin_entry_point(entry_point)
+
+    assert Stream.test_double.__name__ == "stub"
+    Stream.test_double()
+    assert Stream.test_double.__name__ == "test"