Python SDK: Add a decorator to auto-wrap responses in DatastarResponse (#925)

* Add a decorator to auto-wrap responses in DatastarResponse

* Add datastar_response decorator for sanic

* Make the datastar_response generator work the same for sanic as it does for other frameworks

* Explicitly close connection when done with sanic streaming

* Document response decorator in the readme
More readme updates

* Formatting

Author: gazpachoking

README.md

@@ -58,34 +58,54 @@ e.g. `from datastar_py.quart import DatastarResponse` The containing functions
 are not shown here, as they will differ per framework.
 
 ```python
+from datastar_py import ServerSentEventGenerator as SSE
+
 # 0 events, a 204
 return DatastarResponse()
 # 1 event
-return DatastarResponse(ServerSentEventGenerator.patch_elements("<div id='mydiv'></div>"))
+return DatastarResponse(SSE.patch_elements("<div id='mydiv'></div>"))
 # 2 events
 return DatastarResponse([
-    ServerSentEventGenerator.patch_elements("<div id='mydiv'></div>"),
-    ServerSentEventGenerator.patch_signals({"mysignal": "myval"}),
+    SSE.patch_elements("<div id='mydiv'></div>"),
+    SSE.patch_signals({"mysignal": "myval"}),
 ])
 
-
 # N events, a long lived stream (for all frameworks but sanic)
 async def updates():
     while True:
-        yield ServerSentEventGenerator.patch_elements("<div id='mydiv'></div>")
+        yield SSE.patch_elements("<div id='mydiv'></div>")
         await asyncio.sleep(1)
-
-
 return DatastarResponse(updates())
+
 # A long lived stream for sanic
 response = await datastar_respond(request)
 # which is just a helper for the following
 # response = await request.respond(DatastarResponse())
 while True:
-    await response.send(ServerSentEventGenerator.patch_elements("<div id='mydiv'></div>"))
+    await response.send(SSE.patch_elements("<div id='mydiv'></div>"))
     await asyncio.sleep(1)
 ```
 
+### Response Decorator
+To make returning a `DatastarResponse` simpler, there is a decorator
+`datastar_response` available that automatically wraps a function result in
+`DatastarResponse`. It works on async and regular functions and generator
+functions. The main use case is when using a generator function, as you can
+avoid a second generator function inside your response function. The decorator
+works the same for any of the supported frameworks, and should be used under
+any routing decorator from the framework.
+
+```python
+from datastar_py.sanic import datastar_response, ServerSentEventGenerator as SSE
+
+@app.get('/my_route')
+@datastar_response
+def my_route(request):
+    while True:
+        yield SSE.patch_elements("<div id='mydiv'></div>")
+        await asyncio.sleep(1)
+```
+
 ## Signal Helpers
 The current state of the datastar signals is included by default in every 
 datastar request. A helper is included to load those signals for each
@@ -117,4 +137,16 @@ Button("My Button", data.on("click", "console.log('clicked')").debounce(1000).st
 f"<button {data.on("click", "console.log('clicked')").debounce(1000).stop}>My Button</button>"
 # Jinja, but no editor completion :(
 <button {{data.on("click", "console.log('clicked')").debounce(1000).stop}}>My Button</button>
+```
+
+When using datastar with a different alias, you can instantiate the class yourself.
+
+```python
+from datastar_py.attributes import AttributeGenerator
+
+data = AttributeGenerator(alias="data-star-")
+
+# htmy (htmy will transform _ into - unless the attribute starts with _, which will be stripped)
+data = AttributeGenerator(alias="_data-")
+html.button("My Button", **data.on("click", "console.log('clicked')").debounce("1s").stop)
 ```

sdk-test.py

@@ -14,6 +14,7 @@
 2. Move to the sdk/tests folder.
 3. Run `test-all.sh http://127.0.0.1:8000` to run the tests.
 """
+
 import re
 
 from sanic import Request, Sanic

src/datastar_py/attributes.py

@@ -138,7 +138,8 @@ def computed(self, computed_dict: Mapping | None = None, /, **computed: str) ->
         """Create signals that are computed based on an expression."""
         computed = {**(computed_dict if computed_dict else {}), **computed}
         first, *rest = (
-            BaseAttr("computed", key=sig, value=expr, alias=self._alias) for sig, expr in computed.items()
+            BaseAttr("computed", key=sig, value=expr, alias=self._alias)
+            for sig, expr in computed.items()
         )
         first._other_attrs = rest
         return first

src/datastar_py/django.py

@@ -1,17 +1,15 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from collections.abc import Awaitable, Mapping
+from functools import wraps
+from typing import Any, Callable, ParamSpec
 
 from django.http import HttpRequest
 from django.http import StreamingHttpResponse as _StreamingHttpResponse
 
 from . import _read_signals
 from .sse import SSE_HEADERS, DatastarEvent, DatastarEvents, ServerSentEventGenerator
 
-if TYPE_CHECKING:
-    from collections.abc import Mapping
-
-
 __all__ = [
     "SSE_HEADERS",
     "DatastarResponse",
@@ -42,5 +40,26 @@ def __init__(
         super().__init__(content, status=status, headers=headers)
 
 
+P = ParamSpec("P")
+
+
+def datastar_response(
+    func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
+) -> Callable[P, Awaitable[DatastarResponse]]:
+    """A decorator which wraps a function result in DatastarResponse.
+
+    Can be used on a sync or async function or generator function.
+    """
+
+    @wraps(func)
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+        r = func(*args, **kwargs)
+        if isinstance(r, Awaitable):
+            return DatastarResponse(await r)
+        return DatastarResponse(r)
+
+    return wrapper
+
+
 def read_signals(request: HttpRequest) -> dict[str, Any] | None:
     return _read_signals(request.method, request.headers, request.GET, request.body)

src/datastar_py/fastapi.py

@@ -3,13 +3,14 @@
 from fastapi import Depends
 
 from .sse import SSE_HEADERS, ServerSentEventGenerator
-from .starlette import DatastarResponse, read_signals
+from .starlette import DatastarResponse, datastar_response, read_signals
 
 __all__ = [
     "SSE_HEADERS",
     "DatastarResponse",
     "ReadSignals",
     "ServerSentEventGenerator",
+    "datastar_response",
     "read_signals",
 ]
 

src/datastar_py/fasthtml.py

@@ -1,9 +1,10 @@
 from .sse import SSE_HEADERS, ServerSentEventGenerator
-from .starlette import DatastarResponse, read_signals
+from .starlette import DatastarResponse, datastar_response, read_signals
 
 __all__ = [
     "SSE_HEADERS",
     "DatastarResponse",
     "ServerSentEventGenerator",
+    "datastar_response",
     "read_signals",
 ]

src/datastar_py/litestar.py

@@ -1,15 +1,20 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from collections.abc import Awaitable, Mapping
+from functools import wraps
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    ParamSpec,
+)
 
 from litestar.response import Stream
 
 from . import _read_signals
 from .sse import SSE_HEADERS, DatastarEvent, DatastarEvents, ServerSentEventGenerator
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping
-
     from litestar import Request
     from litestar.background_tasks import BackgroundTask, BackgroundTasks
     from litestar.types import ResponseCookies
@@ -54,6 +59,28 @@ def __init__(
         )
 
 
+P = ParamSpec("P")
+
+
+def datastar_response(
+    func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
+) -> Callable[P, Awaitable[DatastarResponse]]:
+    """A decorator which wraps a function result in DatastarResponse.
+
+    Can be used on a sync or async function or generator function.
+    """
+
+    @wraps(func)
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+        r = func(*args, **kwargs)
+        if isinstance(r, Awaitable):
+            return DatastarResponse(await r)
+        return DatastarResponse(r)
+
+    wrapper.__annotations__["return"] = "DatastarResponse"
+    return wrapper
+
+
 async def read_signals(request: Request) -> dict[str, Any] | None:
     return _read_signals(
         request.method, request.headers, request.query_params, await request.body()

src/datastar_py/quart.py

@@ -1,16 +1,15 @@
 from __future__ import annotations
 
-from inspect import isasyncgen, isgenerator
-from typing import TYPE_CHECKING, Any
+from collections.abc import Awaitable, Mapping
+from functools import wraps
+from inspect import isasyncgen, isasyncgenfunction, isgenerator
+from typing import Any, Callable, ParamSpec
 
-from quart import Response, request
+from quart import Response, copy_current_request_context, request, stream_with_context
 
 from . import _read_signals
 from .sse import SSE_HEADERS, DatastarEvents, ServerSentEventGenerator
 
-if TYPE_CHECKING:
-    from collections.abc import Mapping
-
 __all__ = [
     "SSE_HEADERS",
     "DatastarResponse",
@@ -39,5 +38,26 @@ def __init__(
             self.timeout = None
 
 
+P = ParamSpec("P")
+
+
+def datastar_response(
+    func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
+) -> Callable[P, Awaitable[DatastarResponse]]:
+    """A decorator which wraps a function result in DatastarResponse.
+
+    Can be used on a sync or async function or generator function.
+    """
+
+    @wraps(func)
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+        if isasyncgenfunction(func):
+            return DatastarResponse(stream_with_context(func)(*args, **kwargs))
+        return DatastarResponse(await copy_current_request_context(func)(*args, **kwargs))
+
+    wrapper.__annotations__["return"] = "DatastarResponse"
+    return wrapper
+
+
 async def read_signals() -> dict[str, Any] | None:
     return _read_signals(request.method, request.headers, request.args, await request.get_data())

src/datastar_py/sanic.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from collections.abc import Awaitable, Collection, Mapping
+from functools import wraps
+from inspect import isasyncgen, isgenerator
+from typing import Any, Callable, ParamSpec
 
 from sanic import HTTPResponse, Request
 
 from . import _read_signals
-from .sse import SSE_HEADERS, DatastarEvent, ServerSentEventGenerator
-
-if TYPE_CHECKING:
-    from collections.abc import Collection, Mapping
+from .sse import SSE_HEADERS, DatastarEvent, DatastarEvents, ServerSentEventGenerator
 
 __all__ = [
     "SSE_HEADERS",
@@ -52,5 +52,41 @@ async def datastar_respond(
     return await request.respond(DatastarResponse(status=status, headers=headers))
 
 
+P = ParamSpec("P")
+
+
+def datastar_response(
+    func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
+) -> Callable[P, Awaitable[DatastarResponse]]:
+    """A decorator which wraps a function result in DatastarResponse.
+
+    Can be used on a sync or async function or generator function.
+    """
+
+    @wraps(func)
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+        r = func(*args, **kwargs)
+        if isinstance(r, Awaitable):
+            return DatastarResponse(await r)
+        if isasyncgen(r):
+            request = args[0]
+            response = await request.respond(response=DatastarResponse())
+            async for event in r:
+                await response.send(event)
+            await response.eof()
+            return response
+        if isgenerator(r):
+            request = args[0]
+            response = await request.respond(response=DatastarResponse())
+            for event in r:
+                await response.send(event)
+            await response.eof()
+            return response
+        return DatastarResponse(r)
+
+    wrapper.__annotations__["return"] = "DatastarResponse"
+    return wrapper
+
+
 async def read_signals(request: Request) -> dict[str, Any] | None:
     return _read_signals(request.method, request.headers, request.args, request.body)

src/datastar_py/sse.py

@@ -150,8 +150,7 @@ def patch_signals(
             signals if isinstance(signals, str) else json.dumps(signals, separators=(",", ":"))
         )
         data_lines.extend(
-            f"{consts.SIGNALS_DATALINE_LITERAL} {line}"
-            for line in signals_str.splitlines()
+            f"{consts.SIGNALS_DATALINE_LITERAL} {line}" for line in signals_str.splitlines()
         )
 
         return ServerSentEventGenerator._send(