add datastar SSE rendering

repub/datastar.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+from collections.abc import AsyncGenerator, Awaitable, Callable
+from typing import Protocol
+
+from datastar_py import ServerSentEventGenerator as SSE
+from datastar_py.sse import DatastarEvent
+
+
+class HtmlRenderable(Protocol):
+    def __html__(self) -> str: ...
+
+
+RenderResult = str | HtmlRenderable
+RenderFunction = Callable[[], Awaitable[RenderResult]]
+
+
+class RefreshBroker:
+    def __init__(self) -> None:
+        self._subscribers: set[asyncio.Queue[object]] = set()
+
+    def subscribe(self) -> asyncio.Queue[object]:
+        queue: asyncio.Queue[object] = asyncio.Queue(maxsize=1)
+        self._subscribers.add(queue)
+        return queue
+
+    def unsubscribe(self, queue: asyncio.Queue[object]) -> None:
+        self._subscribers.discard(queue)
+
+    def publish(self, event: object = "refresh-event") -> None:
+        for queue in tuple(self._subscribers):
+            if queue.full():
+                try:
+                    queue.get_nowait()
+                except asyncio.QueueEmpty:
+                    pass
+            try:
+                queue.put_nowait(event)
+            except asyncio.QueueFull:
+                continue
+
+
+async def render_sse_event(
+    render: RenderFunction, *, last_event_id: str | None = None
+) -> tuple[str | None, DatastarEvent | None]:
+    html = _coerce_html(await render())
+    event_id = _render_hash(html)
+    if event_id == last_event_id:
+        return last_event_id, None
+    return event_id, SSE.patch_elements(html, event_id=event_id)
+
+
+async def render_stream(
+    queue: asyncio.Queue[object],
+    render: RenderFunction,
+    *,
+    last_event_id: str | None = None,
+    render_on_connect: bool = True,
+) -> AsyncGenerator[DatastarEvent, None]:
+    if render_on_connect:
+        last_event_id, event = await render_sse_event(
+            render, last_event_id=last_event_id
+        )
+        if event is not None:
+            yield event
+
+    while True:
+        await queue.get()
+        last_event_id, event = await render_sse_event(
+            render, last_event_id=last_event_id
+        )
+        if event is not None:
+            yield event
+
+
+def _coerce_html(view: RenderResult) -> str:
+    if isinstance(view, str):
+        return view
+    return view.__html__()
+
+
+def _render_hash(html: str) -> str:
+    return hashlib.blake2s(html.encode("utf-8"), digest_size=16).hexdigest()