41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Iterable
|
|
|
|
from fastapi import WebSocket
|
|
|
|
|
|
class RealtimeHub:
|
|
def __init__(self) -> None:
|
|
self._connections: set[WebSocket] = set()
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|
|
|
def bind_loop(self, loop: asyncio.AbstractEventLoop) -> None:
|
|
self._loop = loop
|
|
|
|
async def connect(self, websocket: WebSocket) -> None:
|
|
await websocket.accept()
|
|
self._connections.add(websocket)
|
|
|
|
def disconnect(self, websocket: WebSocket) -> None:
|
|
self._connections.discard(websocket)
|
|
|
|
async def broadcast(self, payload: dict) -> None:
|
|
stale: list[WebSocket] = []
|
|
for websocket in self._connections:
|
|
try:
|
|
await websocket.send_json(payload)
|
|
except Exception:
|
|
stale.append(websocket)
|
|
for websocket in stale:
|
|
self._connections.discard(websocket)
|
|
|
|
def notify(self, payload: dict) -> None:
|
|
if self._loop is None:
|
|
return
|
|
asyncio.run_coroutine_threadsafe(self.broadcast(payload), self._loop)
|
|
|
|
|
|
hub = RealtimeHub()
|