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.quart import DatastarResponse 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)) 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)