47 lines
1.4 KiB
Python
Executable File
47 lines
1.4 KiB
Python
Executable File
"""WebSocket connection manager for real-time version room events."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from fastapi import WebSocket
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self) -> None:
|
|
# version_id -> list of active WebSocket connections
|
|
self._rooms: dict[str, list[WebSocket]] = {}
|
|
|
|
async def connect(self, version_id: uuid.UUID, websocket: WebSocket) -> None:
|
|
await websocket.accept()
|
|
key = str(version_id)
|
|
self._rooms.setdefault(key, []).append(websocket)
|
|
|
|
def disconnect(self, version_id: uuid.UUID, websocket: WebSocket) -> None:
|
|
key = str(version_id)
|
|
room = self._rooms.get(key, [])
|
|
if websocket in room:
|
|
room.remove(websocket)
|
|
if not room:
|
|
self._rooms.pop(key, None)
|
|
|
|
async def broadcast(self, version_id: uuid.UUID, event: str, data: Any) -> None:
|
|
key = str(version_id)
|
|
payload = json.dumps({"event": event, "data": data})
|
|
dead: list[WebSocket] = []
|
|
for ws in self._rooms.get(key, []):
|
|
try:
|
|
await ws.send_text(payload)
|
|
except Exception:
|
|
dead.append(ws)
|
|
for ws in dead:
|
|
self.disconnect(version_id, ws)
|
|
|
|
def room_size(self, version_id: uuid.UUID) -> int:
|
|
return len(self._rooms.get(str(version_id), []))
|
|
|
|
|
|
manager = ConnectionManager()
|