50 lines
2.0 KiB
Python
50 lines
2.0 KiB
Python
from typing import TYPE_CHECKING, Dict, List
|
|
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):
|
|
if connection.ws.application_state == WebSocketState.DISCONNECTED:
|
|
self.remove(group, connection)
|
|
elif connection.ws.application_state == WebSocketState.CONNECTED:
|
|
await connection.send(message)
|
|
|
|
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)
|
|
|
|
async def broadcast(self, message, group: str, exclude: list["RoomConsumer"] = []):
|
|
print('BROADCaST', message)
|
|
if group in self.active_connections:
|
|
for connection in list(set(self.active_connections[group])):
|
|
# print(connection)
|
|
if connection not in exclude:
|
|
await self._send(connection, message, group)
|
|
|
|
async def send_to(self, group, id_code, msg):
|
|
if group in self.active_connections:
|
|
members = [c for c in self.active_connections[group]
|
|
if c.member.id_code == id_code]
|
|
for m in members:
|
|
await self._send(m, msg, group)
|
|
|
|
async def send_to_admin(self, group, msg):
|
|
if group in self.active_connections:
|
|
members = [c for c in self.active_connections[group]
|
|
if c.member.is_admin == True]
|
|
for m in members:
|
|
await self._send(m, msg, group)
|