Enhancement: Support for returning AsyncGenerator[ServerSentEventMessage, None] at route handler
Summary
Currently, Litestar forces you to create an intermediary function that yields ServerSentEventMessage then wrap the generator in ServerSentEvent.
class SomeController(Controller):
async def intermediary_function(self) -> AsyncGenerator[ServerSentEventMessage, None]:
yield ServerSentEventMessage(data='Doing this..')
yield ServerSentEventMessage(data='Doing that..')
yield ServerSentEventMessage(data='Finished doing this and that..')
@get()
async def actual_server_sent_event_handler(self) -> ServerSentEvent:
return ServerSentEvent(self.intermediary_function()) This is not an ideal pattern as you now have a method that isn't a route handler in your controller. If you were to abstract it out as a function, you'll now have to find an elegant place to define this function. In any case, this is an annoying problem.
Basic Example
Ideally, this is what you'd want. It has good code locality and you don't have to jump through functions or files to understand the full context of the code.
class SomeController(Controller):
@get()
async def actual_server_sent_event_handler(self) -> AsyncGenerator[ServerSentEventMessage, None]:
yield ServerSentEventMessage(data='Doing this..')
yield ServerSentEventMessage(data='Doing that..')
yield ServerSentEventMessage(data='Finished doing this and that..')Drawbacks and Impact
I have a pseudo solution to this using decorators that has zero runtime cost aside from the initial injection at startup. Ideally, the actual solution should not have to depend on any decorators.
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from functools import partial, update_wrapper
from inspect import isasyncgenfunction
from litestar.response import ServerSentEvent, ServerSentEventMessage
async def async_handler[**P](
handler: Callable[P, Iterator[ServerSentEventMessage] | AsyncIterator[ServerSentEventMessage]],
*args: P.args,
**kwargs: P.kwargs,
) -> ServerSentEvent:
return ServerSentEvent(handler(*args, **kwargs))
def server_sent_event[**P](
handler: Callable[P, Iterator[ServerSentEventMessage] | AsyncIterator[ServerSentEventMessage]],
) -> Callable[P, ServerSentEvent | Awaitable[ServerSentEvent]]:
new_handler: Callable[P, ServerSentEvent | Awaitable[ServerSentEvent]] = (
partial(async_handler, handler)
if isasyncgenfunction(handler)
else lambda *args, **kwargs: ServerSentEvent(handler(*args, **kwargs))
)
new_handler = update_wrapper(new_handler, handler)
new_handler.__annotations__['return'] = ServerSentEvent
return new_handler
class SomeController(Controller):
@get()
@server_sent_event
async def actual_server_sent_event_handler(self) -> AsyncGenerator[ServerSentEventMessage, None]:
yield ServerSentEventMessage(data='Doing this..')
yield ServerSentEventMessage(data='Doing that..')
yield ServerSentEventMessage(data='Finished doing this and that..')With this solution, I am still able to correctly generate OpenAPI docs and have no issues with DI. I am quite confident that this can be implemented without issues. If you know of any potential drawbacks, I'd like to hear it.
Unresolved questions
No response
How does funding with Polar work?
Pay now to fund the work behind this issue.
Get updates on progress being made.
Maintainer is rewarded once the issue is completed.
Backer
You're funding impactful open source efforts
Contributor
You want to contribute to this effort
Maintainer
You want to get funding like this too