Clean up some of the magic helpers and add typing

Author: gazpachoking

examples/python/fastapi/app.py

@@ -84,11 +84,14 @@ async def updates2():
 # This is also identical, but yielding a string of fragments automatically calls merge_fragments
 # and dicts automatically calls merge_signals
 @app.get("/updates3")
+# Wraps the resulting async generator in a DatastarStreamingResponse
 @sse_generator
 async def updates3():
     while True:
+        # Implicit merge_fragments
         yield f"""<span id="currentTime">{datetime.now().isoformat()}"""
         await asyncio.sleep(1)
+        # Implicit merge_signals
         yield {"currentTime": f"{datetime.now().isoformat()}"}
         await asyncio.sleep(1)
 

sdk/python/src/datastar_py/quart.py

@@ -2,7 +2,7 @@
 
 from quart import make_response as _make_response
 
-from .sse import ServerSentEventGenerator, SSE_HEADERS, _async_map, _wrap_event
+from .sse import ServerSentEventGenerator, SSE_HEADERS, _sse_iterable_wrapper
 
 
 async def make_datastar_response(async_generator):
@@ -14,7 +14,6 @@ async def make_datastar_response(async_generator):
 def sse_generator(generator_func):
     @wraps(generator_func)
     async def _wrapper(*args, **kwargs):
-        content = _async_map(_wrap_event, generator_func(*args, **kwargs))
-
+        content = _sse_iterable_wrapper(generator_func(*args, **kwargs))
         return await make_datastar_response(content)
     return _wrapper

sdk/python/src/datastar_py/sse.py

@@ -1,6 +1,17 @@
+from __future__ import annotations
+
 import json
 from itertools import chain
-from typing import Optional, Protocol, Union, runtime_checkable
+from typing import (
+    Optional,
+    Protocol,
+    Union,
+    runtime_checkable,
+    Callable,
+    AsyncIterable,
+    TypeVar,
+    Iterable,
+)
 
 import datastar_py.consts as consts
 
@@ -178,15 +189,29 @@ def redirect(cls, location: str):
         return cls.execute_script(f"setTimeout(() => window.location = '{location}')")
 
 
-def _wrap_event(event):
-    if isinstance(event, _HtmlProvider) or (isinstance(event, str) and event.startswith("<")):
+def _wrap_event(event: str | _HtmlProvider | dict) -> str:
+    if isinstance(event, _HtmlProvider) or (
+        isinstance(event, str) and event.startswith("<")
+    ):
         return ServerSentEventGenerator.merge_fragments(event)
     elif isinstance(event, dict):
         return ServerSentEventGenerator.merge_signals(event)
     else:
         return event
 
 
-async def _async_map(func, async_iter):
+async def _async_map(func: Callable, async_iter: AsyncIterable) -> AsyncIterable:
     async for item in async_iter:
         yield func(item)
+
+
+SyncOrAsyncIterable = TypeVar("SyncOrAsyncIterable", AsyncIterable, Iterable)
+
+
+def _sse_iterable_wrapper(iterable: SyncOrAsyncIterable) -> SyncOrAsyncIterable:
+    """Wraps an iterable to allow implicitly turning fragments and dictionaries
+    into merge-fragments and merge-signals events."""
+    if isinstance(iterable, AsyncIterable):
+        return _async_map(_wrap_event, iterable)
+    else:
+        return map(_wrap_event, iterable)

sdk/python/src/datastar_py/starlette.py

@@ -3,7 +3,7 @@
 
 from starlette.responses import StreamingResponse as _StreamingResponse
 
-from .sse import SSE_HEADERS, ServerSentEventGenerator, _wrap_event, _async_map
+from .sse import SSE_HEADERS, ServerSentEventGenerator, _wrap_event, _async_map, _sse_iterable_wrapper
 
 
 class DatastarStreamingResponse(_StreamingResponse):
@@ -17,10 +17,5 @@ def sse_generator(generator_func):
     @wraps(generator_func)
     def _wrapper(*args, **kwargs):
         content = generator_func(*args, **kwargs)
-        if isinstance(content, typing.AsyncIterable):
-            content = _async_map(_wrap_event, content)
-        else:
-            content = map(_wrap_event, content)
-
-        return DatastarStreamingResponse(content)
+        return DatastarStreamingResponse(_sse_iterable_wrapper(content))
     return _wrapper