31 lines
788 B
Python
31 lines
788 B
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from starlette.websockets import WebSocket
|
||
|
|
|
||
|
|
|
||
|
|
class Broadcaster:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self._clients: set[WebSocket] = set()
|
||
|
|
self._lock = asyncio.Lock()
|
||
|
|
|
||
|
|
async def add(self, ws: WebSocket) -> None:
|
||
|
|
async with self._lock:
|
||
|
|
self._clients.add(ws)
|
||
|
|
|
||
|
|
async def remove(self, ws: WebSocket) -> None:
|
||
|
|
async with self._lock:
|
||
|
|
self._clients.discard(ws)
|
||
|
|
|
||
|
|
async def broadcast_json(self, message: dict[str, Any]) -> None:
|
||
|
|
async with self._lock:
|
||
|
|
clients = list(self._clients)
|
||
|
|
for ws in clients:
|
||
|
|
try:
|
||
|
|
await ws.send_json(message)
|
||
|
|
except Exception:
|
||
|
|
await self.remove(ws)
|
||
|
|
|