32 lines
877 B
Python
32 lines
877 B
Python
import json
|
|
from collections import defaultdict
|
|
|
|
from fastapi import WebSocket
|
|
|
|
|
|
class Hub:
|
|
def __init__(self) -> None:
|
|
self.rooms: dict[str, list[WebSocket]] = defaultdict(list)
|
|
|
|
async def join(self, session_id: str, ws: WebSocket) -> None:
|
|
await ws.accept()
|
|
self.rooms[session_id].append(ws)
|
|
|
|
def leave(self, session_id: str, ws: WebSocket) -> None:
|
|
conns = self.rooms.get(session_id, [])
|
|
if ws in conns:
|
|
conns.remove(ws)
|
|
|
|
async def broadcast(self, session_id: str, payload: dict) -> None:
|
|
dead = []
|
|
for ws in list(self.rooms.get(session_id, [])):
|
|
try:
|
|
await ws.send_text(json.dumps(payload))
|
|
except Exception:
|
|
dead.append(ws)
|
|
for ws in dead:
|
|
self.leave(session_id, ws)
|
|
|
|
|
|
hub = Hub()
|