Refine decorator to preserve sync/async semantics

mmacpherson · mmacpherson · commit 598e862505fe · 2025-12-20T14:25:57.000-08:00
diff --git a/src/datastar_py/django.py b/src/datastar_py/django.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Mapping
-from functools import wraps
-from inspect import isawaitable
+from functools import partial, wraps
+from inspect import isasyncgenfunction, iscoroutinefunction
 from typing import Any, Callable, ParamSpec
 
 from django.http import HttpRequest
@@ -46,30 +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 hasattr(r, "__aiter__"):
-            raise NotImplementedError(
-                "Async generators/iterables are not yet supported by the Django adapter; "
-                "use a sync generator or return a single value/awaitable instead."
-            )
-
-        if hasattr(r, "__iter__") and not isinstance(r, (str, bytes)):
-            return DatastarResponse(r)
-
-        if isawaitable(r):
-            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:
diff --git a/src/datastar_py/litestar.py b/src/datastar_py/litestar.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Mapping
-from functools import wraps
-from inspect import isawaitable
+from functools import partial, wraps
+from inspect import isasyncgenfunction, iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -65,32 +65,47 @@ def __init__(
 
 def datastar_response(
     func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
-) -> Callable[P, 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)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
-        r = func(*args, **kwargs)
+    # Case A: Async Generator (async def + yield)
+    if isasyncgenfunction(actual_func):
 
-        if hasattr(r, "__aiter__"):
-            return DatastarResponse(r)
+        @wraps(actual_func)
+        async def async_gen_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            return DatastarResponse(func(*args, **kwargs))
 
-        if hasattr(r, "__iter__") and not isinstance(r, (str, bytes)):
-            return DatastarResponse(r)
+        async_gen_wrapper.__annotations__["return"] = DatastarResponse
+        return async_gen_wrapper
 
-        if isawaitable(r):
-            async def await_and_yield():
-                yield await r
+    # Case B: Standard Coroutine (async def + return)
+    elif iscoroutinefunction(actual_func):
 
-            return DatastarResponse(await_and_yield())
+        @wraps(actual_func)
+        async def async_coro_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            result = await func(*args, **kwargs)
+            return DatastarResponse(result)
 
-        return DatastarResponse(r)
+        async_coro_wrapper.__annotations__["return"] = DatastarResponse
+        return async_coro_wrapper
 
-    wrapper.__annotations__["return"] = DatastarResponse
-    return 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:
diff --git a/src/datastar_py/quart.py b/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:
diff --git a/src/datastar_py/starlette.py b/src/datastar_py/starlette.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Mapping
-from functools import wraps
-from inspect import isawaitable
+from functools import partial, wraps
+from inspect import isasyncgenfunction, iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -55,37 +55,47 @@ def __init__(
 
 def datastar_response(
     func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
-) -> Callable[P, 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)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
-        r = func(*args, **kwargs)
+    # Case A: Async Generator (async def + yield)
+    if isasyncgenfunction(actual_func):
 
-        # Check for async generator/iterator first (most specific case)
-        if hasattr(r, "__aiter__"):
-            return DatastarResponse(r)
+        @wraps(actual_func)
+        async def async_gen_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            return DatastarResponse(func(*args, **kwargs))
 
-        # Check for sync generator/iterator (before Awaitable to avoid false positives)
-        if hasattr(r, "__iter__") and not isinstance(r, (str, bytes)):
-            return DatastarResponse(r)
+        async_gen_wrapper.__annotations__["return"] = DatastarResponse
+        return async_gen_wrapper
 
-        # Check for coroutines/tasks (but NOT async generators, already handled above)
-        if isawaitable(r):
-            # Wrap awaitable in an async generator that yields the result
-            async def await_and_yield():
-                yield await r
+    # Case B: Standard Coroutine (async def + return)
+    elif iscoroutinefunction(actual_func):
 
-            return DatastarResponse(await_and_yield())
+        @wraps(actual_func)
+        async def async_coro_wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
+            result = await func(*args, **kwargs)
+            return DatastarResponse(result)
 
-        # Default case: single value or unknown type
-        return DatastarResponse(r)
+        async_coro_wrapper.__annotations__["return"] = DatastarResponse
+        return async_coro_wrapper
 
-    wrapper.__annotations__["return"] = DatastarResponse
-    return 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:
diff --git a/tests/test_datastar_decorator_runtime.py b/tests/test_datastar_decorator_runtime.py
@@ -35,8 +35,8 @@ def anyio_backend() -> str:
         "async_generator",
     ],
 )
-def test_decorator_returns_response_objects(module_path: str, variant: str) -> None:
-    """Decorated handlers should stay sync-callable and return DatastarResponse immediately."""
+def test_decorator_preserves_sync_async_semantics(module_path: str, variant: str) -> None:
+    """Decorated handlers should preserve sync/async nature of the original function."""
 
     mod = importlib.import_module(module_path)
     datastar_response = mod.datastar_response
@@ -59,12 +59,18 @@ async def handler() -> Any:
         async def handler() -> Any:
             yield SSE.patch_signals({"ok": True})
 
-    result = handler()
-    if inspect.iscoroutine(result):
-        result.close()  # avoid "coroutine was never awaited" warnings
+    is_async_variant = variant.startswith("async_")
 
-    assert not inspect.iscoroutinefunction(handler), "Decorator should preserve sync callable semantics"
-    assert isinstance(result, DatastarResponse)
+    # Verify the wrapper preserves sync/async nature
+    if is_async_variant:
+        assert inspect.iscoroutinefunction(handler), "Async handlers should remain async"
+        # Call and close coroutine to avoid warnings (we can't await in sync test)
+        coro = handler()
+        coro.close()
+    else:
+        assert not inspect.iscoroutinefunction(handler), "Sync handlers should remain sync"
+        result = handler()
+        assert isinstance(result, DatastarResponse), "Sync handlers should return DatastarResponse directly"
 
 
 async def _fetch(
@@ -125,3 +131,54 @@ async def ping(request) -> PlainTextResponse:  # noqa: ANN001
     finally:
         server.should_exit = True
         thread.join(timeout=2)
+
+
+def test_async_generator_iterates_on_event_loop() -> None:
+    """Async generators should iterate on the event loop, not spawn a thread.
+
+    This addresses the concern that a sync wrapper might cause async handlers
+    to run in the threadpool. The wrapper being sync only affects where the
+    generator object is created (trivial); iteration happens based on iterator
+    type - Starlette's StreamingResponse detects __aiter__ and iterates async.
+
+    This test uses Starlette, but the same principle applies to Litestar which
+    also uses a sync wrapper. Litestar's Stream response similarly detects
+    async iterators and iterates them on the event loop.
+    """
+    from starlette.testclient import TestClient
+
+    from datastar_py.starlette import datastar_response
+
+    execution_threads: dict[str, str] = {}
+
+    @datastar_response
+    async def async_gen_handler(request) -> Any:  # noqa: ANN001
+        execution_threads["async_gen"] = threading.current_thread().name
+        yield SSE.patch_signals({"async": True})
+
+    @datastar_response
+    def sync_gen_handler(request) -> Any:  # noqa: ANN001
+        execution_threads["sync_gen"] = threading.current_thread().name
+        yield SSE.patch_signals({"sync": True})
+
+    app = Starlette(routes=[
+        Route("/async", async_gen_handler),
+        Route("/sync", sync_gen_handler),
+    ])
+
+    with TestClient(app) as client:
+        client.get("/async")
+        client.get("/sync")
+
+    # Async generator runs on the asyncio portal thread (event loop context)
+    # Sync generator runs in a separate threadpool worker
+    # The key assertion: they run in DIFFERENT thread contexts
+    assert execution_threads["async_gen"] != execution_threads["sync_gen"], (
+        f"Async and sync generators should run in different thread contexts. "
+        f"Async ran on: {execution_threads['async_gen']}, Sync ran on: {execution_threads['sync_gen']}"
+    )
+
+    # Async generator should be on the event loop thread (asyncio-portal-* or MainThread)
+    assert "asyncio" in execution_threads["async_gen"] or execution_threads["async_gen"] == "MainThread", (
+        f"Async generator should run on event loop, but ran on {execution_threads['async_gen']}"
+    )
