38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
import asyncio
|
|
from typing import Dict, Optional
|
|
from fastapi import WebSocket
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.connections: Dict[str, WebSocket] = {}
|
|
self.lock = asyncio.Lock()
|
|
|
|
async def add_connection(self, device_id: str, websocket: WebSocket):
|
|
async with self.lock:
|
|
self.connections[device_id] = websocket
|
|
|
|
async def remove_connection(self, device_id: str, websocket: Optional[WebSocket] = None):
|
|
async with self.lock:
|
|
if device_id in self.connections and (websocket is None or self.connections[device_id] is websocket):
|
|
del self.connections[device_id]
|
|
|
|
async def get_connection(self, device_id: str) -> Optional[WebSocket]:
|
|
async with self.lock:
|
|
return self.connections.get(device_id)
|
|
|
|
async def is_connected(self, device_id: str) -> bool:
|
|
websocket = await self.get_connection(device_id)
|
|
return (
|
|
websocket is not None
|
|
and getattr(getattr(websocket, "client_state", None), "name", None) == "CONNECTED"
|
|
and not getattr(websocket, "_closed", False)
|
|
)
|
|
|
|
async def get_all_connections(self) -> Dict[str, WebSocket]:
|
|
async with self.lock:
|
|
return dict(self.connections)
|
|
|
|
|
|
connection_manager = ConnectionManager()
|