36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
"""In-process event fan-out for Server-Sent Events clients."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
|
|
class EventBus:
|
|
def __init__(self) -> None:
|
|
self._subscribers: set[asyncio.Queue[dict[str, Any]]] = set()
|
|
|
|
async def publish(self, event_type: str, data: dict[str, Any]) -> None:
|
|
event = {"type": event_type, "data": data}
|
|
dead: list[asyncio.Queue[dict[str, Any]]] = []
|
|
for queue in self._subscribers:
|
|
try:
|
|
queue.put_nowait(event)
|
|
except asyncio.QueueFull:
|
|
dead.append(queue)
|
|
for queue in dead:
|
|
self._subscribers.discard(queue)
|
|
|
|
async def subscribe(self) -> AsyncIterator[dict[str, Any]]:
|
|
queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=100)
|
|
self._subscribers.add(queue)
|
|
try:
|
|
while True:
|
|
try:
|
|
yield await asyncio.wait_for(queue.get(), timeout=20)
|
|
except TimeoutError:
|
|
yield {"type": "heartbeat", "data": {}}
|
|
finally:
|
|
self._subscribers.discard(queue)
|