plugins lazy loading, specifying attribute_name in register_api, tests

Author: roveo

streamz/__init__.py

@@ -5,7 +5,7 @@
 from .sources import *
 from .plugins import load_plugins
 
-load_plugins()
+load_plugins(Stream)
 
 try:
     from .dask import DaskStream, scatter

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
@@ -290,10 +290,21 @@ 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):
+            attribute = entry_point.load()
+            cls.register_api(
+                modifier=modifier, attribute_name=entry_point.name
+            )(attribute)
+            return attribute(*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/plugins.py

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

streamz/tests/test_plugins.py

@@ -0,0 +1,38 @@
+from streamz.sources import Source
+from streamz import Stream
+
+
+class MockEntryPoint:
+
+    def __init__(self, name, cls):
+        self.name = name
+        self.cls = cls
+
+    def load(self):
+        return self.cls
+
+
+def test_register_plugin_entry_point():
+    class test(Stream):
+        pass
+
+    entry_point = MockEntryPoint("test_node", test)
+    Stream.register_plugin_entry_point(entry_point)
+
+    assert Stream.test_node.__name__ == "stub"
+
+    Stream().test_node()
+
+    assert Stream.test_node.__name__ == "test"
+
+
+def test_register_plugin_entry_point_modifier():
+    class test(Source):
+        pass
+
+    entry_point = MockEntryPoint("from_test", test)
+    Stream.register_plugin_entry_point(entry_point, staticmethod)
+
+    Stream.from_test()
+
+    assert Stream.from_test.__self__ is Stream