Add support for startup and shutdown events

bocadillo/api.py

@@ -14,6 +14,7 @@
 
 from .cors import DEFAULT_CORS_CONFIG
 from .error_handlers import ErrorHandler, convert_exception_to_response
+from .events import EventsMixin
 from .exceptions import HTTPError
 from .hooks import HooksMixin
 from .media import Media
@@ -28,7 +29,9 @@
 from .types import ASGIApp, ASGIAppInstance, WSGIApp
 
 
-class API(TemplatesMixin, RoutingMixin, HooksMixin, metaclass=APIMeta):
+class API(
+    TemplatesMixin, RoutingMixin, HooksMixin, EventsMixin, metaclass=APIMeta
+):
     """The all-mighty API class.
 
     This class implements the [ASGI](https://asgi.readthedocs.io) protocol.
@@ -338,6 +341,9 @@ def find_app(self, scope: dict) -> ASGIAppInstance:
             An ASGI application instance
             (either `self` or an instance of a sub-app).
         """
+        if scope["type"] == "lifespan":
+            return self.handle_lifespan(scope)
+
         path: str = scope["path"]
 
         # Return a sub-mounted extra app, if found

bocadillo/events.py

@@ -0,0 +1,63 @@
+from typing import Callable, List, Tuple, Optional
+
+from starlette.middleware.lifespan import LifespanMiddleware
+
+from .types import ASGIApp, ASGIAppInstance
+
+EventHandler = Callable[[], None]
+
+
+class EventsMixin:
+    """Provide server event handling capabilities."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._events: List[Tuple[str, EventHandler]] = []
+
+    def on(self, event: str, handler: Optional[EventHandler] = None):
+        """Register an event handler.
+
+        # Parameters
+        event (str):
+            Either "startup" (when the server boots) or "shutdown" (when the
+            server stops).
+        handler (callback, optional):
+            The event handler. If not given, this should be used as a
+            decorator.
+        """
+
+        if handler is not None:
+            self._add_event_handler(event, handler)
+            return
+        else:
+
+            def decorate(f: EventHandler):
+                self._add_event_handler(event, f)
+                return f
+
+            return decorate
+
+    def _add_event_handler(self, event: str, handler: EventHandler):
+        self._events.append((event, handler))
+
+    def _get_lifespan_middleware(self, app: ASGIApp):
+        middleware = LifespanMiddleware(app)
+        for event, func in self._events:
+            middleware.add_event_handler(event, func)
+        return middleware
+
+    def handle_lifespan(self, scope: dict) -> ASGIAppInstance:
+        # Strict implementation of the ASGI lifespan spec.
+        # This is required because the Starlette `LifespanMiddleware`
+        # does not send the `complete` responses.
+
+        async def app(receive, send):
+            message = await receive()
+            assert message["type"] == "lifespan.startup"
+            await send({"type": "lifespan.startup.complete"})
+
+            message = await receive()
+            assert message["type"] == "lifespan.shutdown"
+            await send({"type": "lifespan.shutdown.complete"})
+
+        return self._get_lifespan_middleware(lambda s: app)(scope)

docs/.vuepress/config.js

@@ -97,6 +97,7 @@ module.exports = {
                         '/topics/features/cors',
                         '/topics/features/hsts',
                         '/topics/features/gzip',
+                        '/topics/features/events',
                         '/topics/features/middleware',
                     ],
                 },

docs/api/api.md

@@ -117,6 +117,21 @@ __Example__
 ...     pass
 ```
 
+## on
+```python
+API.on(self, event: str, handler: Union[Callable[[], NoneType], NoneType] = None)
+```
+Register an event handler.
+
+__Parameters__
+
+- __event (str)__:
+    Either "startup" (when the server boots) or "shutdown" (when the
+    server stops).
+- __handler (callback, optional)__:
+    The event handler. If not given, this should be used as a
+    decorator.
+
 ## url_for
 ```python
 API.url_for(self, name: str, **kwargs) -> str

docs/topics/features/events.md

@@ -0,0 +1,38 @@
+# Events
+
+Bocadillo implements the [ASGI Lifespan protocol](https://asgi.readthedocs.io/en/latest/specs/lifespan.html), which allows you to hook into the server's initialization and teardown via events and event handlers.
+
+Events are especially helpful when you need to setup resources on server startup and make sure you clean them up when the server stops.
+
+## Registering events
+
+Event handlers are callbacks, i.e. functions with the signature `() -> None`.
+
+They can be registered using the `@api.on()` decorator:
+
+```python
+@api.on("startup")
+async def setup():
+    # Perform setup when server boots
+    pass
+
+@api.on("shutdown")
+async def cleanup():
+    # Perform cleanup when server shuts down
+    pass
+```
+
+A non-decorator syntax is also available:
+
+```python
+async def setup():
+    pass
+
+api.on("shutdown", setup)
+```
+
+Only the `"startup"` and `"shutdown"` events are supported.
+
+::: tip
+Event handlers can also be regular, non-async functions.
+:::

tests/test_events.py

@@ -0,0 +1,52 @@
+from bocadillo import API
+
+
+def test_startup_and_shutdown(api: API):
+    message = None
+
+    @api.on("startup")
+    async def setup():
+        nonlocal message
+        message = "hi"
+
+    @api.on("shutdown")
+    async def cleanup():
+        nonlocal message
+        message = None
+
+    @api.route("/")
+    async def index(req, res):
+        res.text = message
+
+    # The Starlette TestClient calls startup and shutdown events when
+    # used as a context manager.
+    with api.client:
+        assert message == "hi"
+        response = api.client.get("/")
+        assert response.text == "hi"
+    assert message is None
+
+
+def test_sync_handler(api: API):
+    message = None
+
+    @api.on("startup")
+    def setup():
+        nonlocal message
+        message = "hi"
+
+    with api.client:
+        assert message == "hi"
+
+
+def test_non_decorator_syntax(api: API):
+    message = None
+
+    async def setup():
+        nonlocal message
+        message = "hi"
+
+    api.on("startup", setup)
+
+    with api.client:
+        assert message == "hi"