from __future__ import annotations import asyncio import hashlib from collections.abc import AsyncGenerator from contextlib import suppress from typing import cast import htpy as h from datastar_py import ServerSentEventGenerator as SSE from datastar_py.quart import DatastarResponse, read_signals from datastar_py.sse import DatastarEvent from htpy import Renderable from quart import Quart, Response, request, url_for from repub.datastar import RefreshBroker, render_stream from repub.pages import admin_component, shim_page REFRESH_BROKER_KEY = "repub.refresh_broker" ACTIVE_JOBS_KEY = "repub.demo_active_jobs" REFRESH_TASK_KEY = "repub.demo_refresh_task" def _render_shim_page(*, stylesheet_href: str, datastar_src: str) -> tuple[str, str]: head = ( h.title["Republisher Admin UI"], h.link(rel="stylesheet", href=stylesheet_href), ) body = str(shim_page(datastar_src=datastar_src, head=head)) etag = hashlib.sha256(body.encode("utf-8")).hexdigest() return body, etag def create_app(*, enable_demo_refresh: bool = True) -> Quart: app = Quart(__name__) app.extensions[REFRESH_BROKER_KEY] = RefreshBroker() app.extensions[ACTIVE_JOBS_KEY] = 12 if enable_demo_refresh: @app.before_serving async def start_demo_refresh() -> None: app.extensions[REFRESH_TASK_KEY] = asyncio.create_task( _demo_refresh_loop(app) ) @app.after_serving async def stop_demo_refresh() -> None: task = cast(asyncio.Task[None] | None, app.extensions.get(REFRESH_TASK_KEY)) if task is None: return task.cancel() with suppress(asyncio.CancelledError): await task @app.get("/") async def index() -> Response: body, etag = _render_shim_page( stylesheet_href=url_for("static", filename="app.css"), datastar_src=url_for("static", filename="datastar@1.0.0-RC.8.js"), ) if request.if_none_match.contains(etag): response = Response(status=304) response.set_etag(etag) return response response = Response(body, mimetype="text/html") response.set_etag(etag) return response @app.post("/") async def index_patch() -> DatastarResponse: queue = get_refresh_broker(app).subscribe() stream = render_stream( queue, render=lambda: render_dashboard(app), last_event_id=request.headers.get("last-event-id"), ) return DatastarResponse(_unsubscribe_on_close(queue, stream, app)) @app.post("/demo/decrement") async def demo_decrement() -> DatastarResponse: amount, error = _validated_decrement_amount(await read_signals()) if error is not None: return DatastarResponse(SSE.patch_signals({"decrementError": error})) set_active_jobs(app, max(0, get_active_jobs(app) - amount)) trigger_refresh(app) return DatastarResponse(SSE.patch_signals({"decrementError": ""})) return app def get_refresh_broker(app: Quart) -> RefreshBroker: return cast(RefreshBroker, app.extensions[REFRESH_BROKER_KEY]) def trigger_refresh(app: Quart, event: object = "refresh-event") -> None: get_refresh_broker(app).publish(event) async def render_dashboard(app: Quart) -> Renderable: return admin_component(active_jobs=str(get_active_jobs(app))) async def _unsubscribe_on_close( queue: object, stream: AsyncGenerator[DatastarEvent, None], app: Quart ) -> AsyncGenerator[DatastarEvent, None]: try: async for event in stream: yield event finally: get_refresh_broker(app).unsubscribe(cast(asyncio.Queue[object], queue)) def get_active_jobs(app: Quart) -> int: return cast(int, app.extensions[ACTIVE_JOBS_KEY]) def set_active_jobs(app: Quart, value: int) -> None: app.extensions[ACTIVE_JOBS_KEY] = value async def _demo_refresh_loop(app: Quart) -> None: while True: await asyncio.sleep(1) set_active_jobs(app, get_active_jobs(app) + 1) trigger_refresh(app) def _validated_decrement_amount( signals: dict[str, object] | None, ) -> tuple[int, str | None]: raw_amount = ( "" if signals is None else str(signals.get("decrementAmount", "")).strip() ) try: amount = int(raw_amount) except ValueError: return 0, "Decrement amount must be an odd integer." if amount < 1 or amount % 2 == 0: return 0, "Decrement amount must be an odd integer." return amount, None