Handle generator and awaitable responses across adapters

Author: mmacpherson

.gitignore

@@ -2,3 +2,4 @@
 .idea
 .zed
 __pycache__
+uv.lock

src/datastar_py/django.py

@@ -54,8 +54,18 @@ def datastar_response(
     @wraps(func)
     async def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
         r = func(*args, **kwargs)
+
+        # Collect async iterables so we don't hand StreamingHttpResponse an async generator
+        if hasattr(r, "__aiter__"):
+            events = [event async for event in r]
+            return DatastarResponse(events)
+
+        if hasattr(r, "__iter__") and not isinstance(r, (str, bytes)):
+            return DatastarResponse(r)
+
         if isinstance(r, Awaitable):
             return DatastarResponse(await r)
+
         return DatastarResponse(r)
 
     return wrapper

src/datastar_py/litestar.py

@@ -64,17 +64,28 @@ def __init__(
 
 def datastar_response(
     func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
-) -> Callable[P, Awaitable[DatastarResponse]]:
+) -> Callable[P, 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:
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
         r = func(*args, **kwargs)
+
+        if hasattr(r, "__aiter__"):
+            return DatastarResponse(r)
+
+        if hasattr(r, "__iter__") and not isinstance(r, (str, bytes)):
+            return DatastarResponse(r)
+
         if isinstance(r, Awaitable):
-            return DatastarResponse(await r)
+            async def await_and_yield():
+                yield await r
+
+            return DatastarResponse(await_and_yield())
+
         return DatastarResponse(r)
 
     wrapper.__annotations__["return"] = DatastarResponse

src/datastar_py/quart.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Awaitable, Mapping
 from functools import wraps
-from inspect import isasyncgen, isasyncgenfunction, isgenerator
+from inspect import isasyncgen, isgenerator
 from typing import Any, Callable, ParamSpec
 
 from quart import Response, copy_current_request_context, request, stream_with_context
@@ -51,9 +51,23 @@ def datastar_response(
 
     @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))
+        # Preserve request context for whatever we return
+        bound_func = copy_current_request_context(func)
+        r = bound_func(*args, **kwargs)
+
+        if hasattr(r, "__aiter__"):
+            return DatastarResponse(stream_with_context(r))
+
+        if hasattr(r, "__iter__") and not isinstance(r, (str, bytes)):
+            return DatastarResponse(stream_with_context(r))
+
+        if isinstance(r, Awaitable):
+            async def await_and_yield():
+                yield await r
+
+            return DatastarResponse(stream_with_context(await_and_yield()))
+
+        return DatastarResponse(r)
 
     wrapper.__annotations__["return"] = DatastarResponse
     return wrapper

src/datastar_py/starlette.py

@@ -54,17 +54,33 @@ def __init__(
 
 def datastar_response(
     func: Callable[P, Awaitable[DatastarEvents] | DatastarEvents],
-) -> Callable[P, Awaitable[DatastarResponse]]:
+) -> Callable[P, 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:
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> DatastarResponse:
         r = func(*args, **kwargs)
+
+        # Check for async generator/iterator first (most specific case)
+        if hasattr(r, "__aiter__"):
+            return DatastarResponse(r)
+
+        # Check for sync generator/iterator (before Awaitable to avoid false positives)
+        if hasattr(r, "__iter__") and not isinstance(r, (str, bytes)):
+            return DatastarResponse(r)
+
+        # Check for coroutines/tasks (but NOT async generators, already handled above)
         if isinstance(r, Awaitable):
-            return DatastarResponse(await r)
+            # Wrap awaitable in an async generator that yields the result
+            async def await_and_yield():
+                yield await r
+
+            return DatastarResponse(await_and_yield())
+
+        # Default case: single value or unknown type
         return DatastarResponse(r)
 
     wrapper.__annotations__["return"] = DatastarResponse