49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
"""Websocket functionality."""
|
|
|
|
import logging
|
|
|
|
from fastapi import WebSocket
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class WebsocketManager:
|
|
"""Websocket manager."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize."""
|
|
self.connections: dict[str, set[WebSocket]] = {}
|
|
|
|
def connect(self, pool: tuple[str, ...], websocket: WebSocket) -> None:
|
|
"""Connect a websocket."""
|
|
pool_str = ",".join(pool)
|
|
if pool_str not in self.connections:
|
|
self.connections[pool_str] = set()
|
|
self.connections[pool_str].add(websocket)
|
|
logger.info(
|
|
"Client %s is now listening on pool %s",
|
|
websocket.client.host if websocket.client else "?",
|
|
pool_str,
|
|
)
|
|
|
|
def disconnect(self, pool: tuple[str, ...], websocket: WebSocket) -> None:
|
|
"""Disconnect a websocket."""
|
|
pool_str = ",".join(pool)
|
|
if pool_str in self.connections:
|
|
self.connections[pool_str].remove(websocket)
|
|
logger.info(
|
|
"Client %s has stopped listening on pool %s",
|
|
websocket.client.host if websocket.client else "?",
|
|
pool_str,
|
|
)
|
|
|
|
async def send(self, pool: tuple[str, ...], message: str) -> None:
|
|
"""Send a message to all websockets in a pool."""
|
|
pool_str = ",".join(pool)
|
|
if pool_str in self.connections:
|
|
for websocket in self.connections[pool_str]:
|
|
await websocket.send_text(message)
|
|
|
|
|
|
websocket_manager = WebsocketManager()
|