align datastar_response behavior with docs (tests, fixes, examples)

Makefile

@@ -0,0 +1,4 @@
+.PHONY: test
+
+test:
+	uv run --dev pytest

README.md

@@ -85,6 +85,10 @@ headers of the SSE response yourself.
 
 A datastar response consists of 0..N datastar events. There are response
 classes included to make this easy in all of the supported frameworks.
+Each framework also exposes a `@datastar_response` decorator that will wrap
+return values (including generators) into the right response class while
+preserving sync handlers as sync so frameworks can keep them in their
+threadpools.
 
 The following examples will work across all supported frameworks when the
 response class is imported from the appropriate framework package.

pyproject.toml

@@ -50,12 +50,14 @@ urls.GitHub = "https://github.com/starfederation/datastar-python"
 dev = [
   "django>=4.2.23",
   "fastapi>=0.116.1",
+  "httpx>=0.27",
   "litestar>=2.17",
   "pre-commit>=4.2",
   "python-fasthtml>=0.12.25; python_full_version>='3.10'",
   "quart>=0.20",
   "sanic>=25.3",
   "starlette>=0.47.3",
+  "uvicorn>=0.30",
 ]
 
 [tool.ruff]

src/datastar_py/django.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Mapping
-from functools import wraps
+from functools import partial, wraps
+from inspect import isasyncgenfunction, iscoroutinefunction
 from typing import Any, Callable, ParamSpec
 
 from django.http import HttpRequest
@@ -45,20 +46,44 @@ def __init__(
 
 def datastar_response(
     func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
-) -> Callable[P, Awaitable[DatastarResponse]]:
+) -> Callable[P, Awaitable[DatastarResponse] | DatastarResponse]:
     """A decorator which wraps a function result in DatastarResponse.
 
     Can be used on a sync or async function or generator function.
+    Preserves the sync/async nature of the decorated 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
+    # Unwrap partials to inspect the actual underlying function
+    actual_func = func
+    while isinstance(actual_func, partial):
+        actual_func = actual_func.func
+
+    # Async generators not supported by Django
+    if isasyncgenfunction(actual_func):
+        raise NotImplementedError(
+            "Async generators are not yet supported by the Django adapter; "
+            "use a sync generator or return a single value/awaitable instead."
+        )
+
+    # Coroutine (async def + return)
+    if iscoroutinefunction(actual_func):
+
+        @wraps(actual_func)
+        async def async_coro_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            result = await func(*args, **kwargs)
+            return DatastarResponse(result)
+
+        async_coro_wrapper.__annotations__["return"] = DatastarResponse
+        return async_coro_wrapper
+
+    # Sync Function (def) - includes sync generators
+    else:
+
+        @wraps(actual_func)
+        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            return DatastarResponse(func(*args, **kwargs))
+
+        sync_wrapper.__annotations__["return"] = DatastarResponse
+        return sync_wrapper
 
 
 def read_signals(request: HttpRequest) -> dict[str, Any] | None:

src/datastar_py/litestar.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Mapping
-from functools import wraps
+from functools import partial, wraps
+from inspect import isasyncgenfunction, iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -64,21 +65,47 @@ def __init__(
 
 def datastar_response(
     func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
-) -> Callable[P, Awaitable[DatastarResponse]]:
+) -> Callable[P, Awaitable[DatastarResponse] | DatastarResponse]:
     """A decorator which wraps a function result in DatastarResponse.
 
     Can be used on a sync or async function or generator function.
+    Preserves the sync/async nature of the decorated function.
     """
+    # Unwrap partials to inspect the actual underlying function
+    actual_func = func
+    while isinstance(actual_func, partial):
+        actual_func = actual_func.func
 
-    @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)
+    # Case A: Async Generator (async def + yield)
+    if isasyncgenfunction(actual_func):
 
-    wrapper.__annotations__["return"] = DatastarResponse
-    return wrapper
+        @wraps(actual_func)
+        async def async_gen_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            return DatastarResponse(func(*args, **kwargs))
+
+        async_gen_wrapper.__annotations__["return"] = DatastarResponse
+        return async_gen_wrapper
+
+    # Case B: Standard Coroutine (async def + return)
+    elif iscoroutinefunction(actual_func):
+
+        @wraps(actual_func)
+        async def async_coro_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            result = await func(*args, **kwargs)
+            return DatastarResponse(result)
+
+        async_coro_wrapper.__annotations__["return"] = DatastarResponse
+        return async_coro_wrapper
+
+    # Case C: Sync Function (def) - includes sync generators
+    else:
+
+        @wraps(actual_func)
+        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            return DatastarResponse(func(*args, **kwargs))
+
+        sync_wrapper.__annotations__["return"] = DatastarResponse
+        return sync_wrapper
 
 
 async def read_signals(request: Request) -> dict[str, Any] | None:

src/datastar_py/quart.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Mapping
-from functools import wraps
-from inspect import isasyncgen, isasyncgenfunction, isgenerator
+from functools import partial, wraps
+from inspect import isasyncgen, isasyncgenfunction, iscoroutinefunction, isgenerator
 from typing import Any, Callable, ParamSpec
 
-from quart import Response, copy_current_request_context, request, stream_with_context
+from quart import Response, request, stream_with_context
 
 from . import _read_signals
 from .sse import SSE_HEADERS, DatastarEvents, ServerSentEventGenerator
@@ -43,20 +43,47 @@ def __init__(
 
 def datastar_response(
     func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
-) -> Callable[P, Awaitable[DatastarResponse]]:
+) -> Callable[P, Awaitable[DatastarResponse] | DatastarResponse]:
     """A decorator which wraps a function result in DatastarResponse.
 
     Can be used on a sync or async function or generator function.
+    Preserves the sync/async nature of the decorated function.
     """
+    # Unwrap partials to inspect the actual underlying function
+    actual_func = func
+    while isinstance(actual_func, partial):
+        actual_func = actual_func.func
 
-    @wraps(func)
-    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
-        if isasyncgenfunction(func):
+    # Case A: Async Generator (async def + yield)
+    if isasyncgenfunction(actual_func):
+
+        @wraps(actual_func)
+        async def async_gen_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
             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_gen_wrapper.__annotations__["return"] = DatastarResponse
+        return async_gen_wrapper
+
+    # Case B: Standard Coroutine (async def + return)
+    elif iscoroutinefunction(actual_func):
+
+        @wraps(actual_func)
+        async def async_coro_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            result = await func(*args, **kwargs)
+            return DatastarResponse(result)
+
+        async_coro_wrapper.__annotations__["return"] = DatastarResponse
+        return async_coro_wrapper
+
+    # Case C: Sync Function (def) - includes sync generators
+    else:
+
+        @wraps(actual_func)
+        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            return DatastarResponse(func(*args, **kwargs))
+
+        sync_wrapper.__annotations__["return"] = DatastarResponse
+        return sync_wrapper
 
 
 async def read_signals() -> dict[str, Any] | None:

src/datastar_py/sanic.py

@@ -4,6 +4,7 @@
 from contextlib import aclosing, closing
 from functools import wraps
 from inspect import isasyncgen, isgenerator
+from inspect import isawaitable
 from typing import Any, Callable, ParamSpec, Union
 
 from sanic import HTTPResponse, Request
@@ -70,7 +71,7 @@ def datastar_response(
     @wraps(func)
     async def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse | None:
         r = func(*args, **kwargs)
-        if isinstance(r, Awaitable):
+        if isawaitable(r):
             return DatastarResponse(await r)
         if isasyncgen(r):
             request = args[0]

src/datastar_py/starlette.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Mapping
-from functools import wraps
+from functools import partial, wraps
+from inspect import isasyncgenfunction, iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -54,21 +55,47 @@ def __init__(
 
 def datastar_response(
     func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
-) -> Callable[P, Awaitable[DatastarResponse]]:
+) -> Callable[P, Awaitable[DatastarResponse] | DatastarResponse]:
     """A decorator which wraps a function result in DatastarResponse.
 
     Can be used on a sync or async function or generator function.
+    Preserves the sync/async nature of the decorated function.
     """
+    # Unwrap partials to inspect the actual underlying function
+    actual_func = func
+    while isinstance(actual_func, partial):
+        actual_func = actual_func.func
 
-    @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)
+    # Case A: Async Generator (async def + yield)
+    if isasyncgenfunction(actual_func):
 
-    wrapper.__annotations__["return"] = DatastarResponse
-    return wrapper
+        @wraps(actual_func)
+        async def async_gen_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            return DatastarResponse(func(*args, **kwargs))
+
+        async_gen_wrapper.__annotations__["return"] = DatastarResponse
+        return async_gen_wrapper
+
+    # Case B: Standard Coroutine (async def + return)
+    elif iscoroutinefunction(actual_func):
+
+        @wraps(actual_func)
+        async def async_coro_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            result = await func(*args, **kwargs)
+            return DatastarResponse(result)
+
+        async_coro_wrapper.__annotations__["return"] = DatastarResponse
+        return async_coro_wrapper
+
+    # Case C: Sync Function (def) - includes sync generators
+    else:
+
+        @wraps(actual_func)
+        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            return DatastarResponse(func(*args, **kwargs))
+
+        sync_wrapper.__annotations__["return"] = DatastarResponse
+        return sync_wrapper
 
 
 async def read_signals(request: Request) -> dict[str, Any] | None: