49 lines
1.5 KiB
Python
Executable file
49 lines
1.5 KiB
Python
Executable file
from fastapi import WebSocket
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manages WebSocket connections for real-time updates."""
|
|
|
|
def __init__(self):
|
|
"""Initialize the connection manager."""
|
|
self.active_connections: dict[str, WebSocket] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, client_id: str) -> None:
|
|
"""
|
|
Accept a new WebSocket connection.
|
|
|
|
Args:
|
|
websocket: The WebSocket connection
|
|
client_id: Unique identifier for this client
|
|
"""
|
|
await websocket.accept()
|
|
self.active_connections[client_id] = websocket
|
|
|
|
def disconnect(self, client_id: str) -> None:
|
|
"""
|
|
Remove a client connection.
|
|
|
|
Args:
|
|
client_id: The client to disconnect
|
|
"""
|
|
if client_id in self.active_connections:
|
|
del self.active_connections[client_id]
|
|
|
|
async def send_message(self, client_id: str, message: dict) -> None:
|
|
"""
|
|
Send a JSON message to a specific client.
|
|
|
|
Args:
|
|
client_id: The target client
|
|
message: Dictionary to send as JSON
|
|
"""
|
|
if client_id in self.active_connections:
|
|
try:
|
|
await self.active_connections[client_id].send_json(message)
|
|
except Exception:
|
|
# Client may have disconnected
|
|
self.disconnect(client_id)
|
|
|
|
def is_connected(self, client_id: str) -> bool:
|
|
"""Check if a client is still connected."""
|
|
return client_id in self.active_connections
|