Add autospawn decorator

biloba/__init__.py

@@ -1,4 +1,5 @@
-from .service import Service, ConfigurableService
+from .service import *  # noqa
+from .service import __all__ as service_all
 from .config import parse_address
 from .util import waitany
 
@@ -10,8 +11,10 @@
 
 
 __all__ = [
-    'ConfigurableService',
-    'Service',
     'parse_address',
     'waitany',
 ]
+
+__all__.extend(service_all)
+
+del service_all

biloba/service.py

@@ -5,6 +5,12 @@
 
 from . import config as biloba_config, events
 
+__all__ = [
+    'Service',
+    'ConfigurableService',
+    'autospawn',
+]
+
 
 class Service(events.EventEmitter):
     """
@@ -291,6 +297,14 @@ def watch_service(self, child):
         finally:
             self.stop()
 
+    def __enter__(self):
+        self.start()
+
+        return self
+
+    def __exit__(self, *exc_args):
+        self.stop()
+
 
 class ConfigurableService(Service):
     """
@@ -328,3 +342,19 @@ def apply_default_config(self, config):
             config.setdefault(key, value)
 
         return config
+
+
+def autospawn(func):
+    """
+    Decorator that will spawn the call in a service's greenlet pool
+
+    The function will only be called once the service has started.
+    """
+    @functools.wraps(func)
+    def wrapped(self, *args, **kwargs):
+        if self.started:
+            self.spawn(func, self, *args, **kwargs)
+        else:
+            self.once('start', lambda: self.spawn(self, *args, **kwargs))
+
+    return wrapped

biloba/test/test_service.py

@@ -462,3 +462,53 @@ def test_biloba_config(self):
         svc = service.ConfigurableService(cfg)
 
         self.assertIs(svc.config, cfg)
+
+
+class AutospawnTestCase(unittest.TestCase):
+    """
+    Tests for `service.autospawn`.
+    """
+
+    def test_started(self):
+        class MyService(service.Service):
+            @service.autospawn
+            def my_func(self):
+                self.emit('autospawn')
+
+        my_service = MyService()
+        self.executed = False
+
+        @my_service.on('autospawn')
+        def on_autospawn():
+            self.assertTrue(my_service.started)
+            self.executed = True
+
+        with my_service:
+            ret = my_service.my_func()
+
+            self.assertIsNone(ret)
+
+        self.assertTrue(self.executed)
+
+    def test_not_started(self):
+        class MyService(service.Service):
+            @service.autospawn
+            def my_func(self):
+                self.emit('autospawn')
+
+        my_service = MyService()
+        self.executed = False
+
+        @my_service.on('autospawn')
+        def on_autospawn():
+            self.assertTrue(my_service.started)
+            self.executed = True
+
+        ret = my_service.my_func()
+
+        self.assertIsNone(ret)
+
+        self.assertFalse(self.executed)
+
+        with my_service:
+            self.assertTrue(self.executed)