-
Notifications
You must be signed in to change notification settings - Fork 0
/
conn_manager.py
48 lines (38 loc) · 1.78 KB
/
conn_manager.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
from collections.abc import Iterable
from fastapi import WebSocket
class ConnectionManager:
def __init__(self):
self._active_connections: set[WebSocket] = set()
self._ws_to_game: dict[WebSocket, str] = dict()
self._game_to_ws: dict[str, list[WebSocket]] = dict()
async def connect(self, websocket: WebSocket):
await websocket.accept()
print(f'Connected {websocket.client}. Active connections count: {len(self._active_connections)}.')
self._active_connections.add(websocket)
def _remove_ws_from_game(self, websocket: WebSocket):
game = self._ws_to_game.get(websocket)
if game:
self._game_to_ws[game].remove(websocket)
self._ws_to_game.pop(websocket)
if len(self._game_to_ws[game]) == 0:
self._game_to_ws.pop(game)
def set_game_name(self, websocket: WebSocket, game_name: str):
if self._ws_to_game.get(websocket) == game_name:
return
self._remove_ws_from_game(websocket)
self._ws_to_game[websocket] = game_name
if game_name not in self._game_to_ws:
self._game_to_ws[game_name] = []
self._game_to_ws[game_name].append(websocket)
def disconnect(self, websocket: WebSocket):
self._active_connections.remove(websocket)
self._remove_ws_from_game(websocket)
print(f'Disconnected {websocket.client}. Active connections count: {len(self._active_connections)}.')
print(self._ws_to_game)
print(self._game_to_ws)
async def send_to_game(self, data: any, game_name: str):
await self.send_json(data, self._game_to_ws.get(game_name, []))
@staticmethod
async def send_json(data: any, ws: Iterable[WebSocket]):
for s in ws:
await s.send_json(data)