from typing import TYPE_CHECKING, Dict, List, Callable, Any from starlette.websockets import WebSocketState if TYPE_CHECKING: from routes.room.consumer import RoomConsumer class RoomManager: def __init__(self): self.active_connections: Dict[str, List["RoomConsumer"]] = {} def add(self, group: str, member: "RoomConsumer"): if group not in self.active_connections: self.active_connections[group] = [] if member not in self.active_connections[group]: self.active_connections[group].append(member) async def _send(self, connection: "RoomConsumer", message, group: str): print('SENDING', connection.ws.application_state, connection.ws.client_state) if connection.ws.application_state == WebSocketState.DISCONNECTED or connection.ws.client_state == WebSocketState.DISCONNECTED: self.remove(group, connection) elif connection.ws.application_state == WebSocketState.CONNECTED and connection.ws.client_state == WebSocketState.CONNECTED: try: await connection.send(message) except: pass def remove(self, group: str, member: "RoomConsumer"): if group in self.active_connections: if member in self.active_connections[group]: self.active_connections[group].remove(member) return True async def broadcast(self, message: Any | Callable, group: str, conditions: list[Callable] = [], exclude: list["RoomConsumer"] = [], ): print('BROADCaST', self.active_connections) if group in self.active_connections: for connection in list(set(self.active_connections[group])): print(connection, connection.ws.state, connection.ws.client_state, connection.ws.application_state) if connection not in exclude and all(f(connection) for f in conditions): await self._send(connection, message, group) async def send_to(self, group, id_code, msg): print('SENDING TO') if group in self.active_connections: members = [c for c in self.active_connections[group] if c.member.id_code == id_code] print('MEM', members) for m in members: await self._send(m, msg, group) async def send_to_admin(self, group, msg): if group in self.active_connections: print('MEMBERS', self.active_connections[group], [(c.ws.state, c.ws.client_state, c.ws.application_state) for c in self.active_connections[group]]) members = [c for c in self.active_connections[group] if c.member is not None and c.member.is_admin is True] for m in members: await self._send(m, msg, group)