diff --git a/bocadillo/api.py b/bocadillo/api.py index 6eb11789..28b6bb00 100644 --- a/bocadillo/api.py +++ b/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 diff --git a/bocadillo/events.py b/bocadillo/events.py new file mode 100644 index 00000000..11148a76 --- /dev/null +++ b/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) diff --git a/docs/.vuepress/config.js b/docs/.vuepress/config.js index 8b5e024c..f342bf31 100644 --- a/docs/.vuepress/config.js +++ b/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', ], }, diff --git a/docs/api/api.md b/docs/api/api.md index 4d9298a6..3815880c 100644 --- a/docs/api/api.md +++ b/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 diff --git a/docs/topics/features/events.md b/docs/topics/features/events.md new file mode 100644 index 00000000..5617015a --- /dev/null +++ b/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. +::: diff --git a/tests/test_events.py b/tests/test_events.py new file mode 100644 index 00000000..5cf159b0 --- /dev/null +++ b/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"