46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
from fastapi import WebSocket
|
|
|
|
|
|
class SocketManager:
|
|
def __init__(self):
|
|
self.rooms = {}
|
|
|
|
async def add_to_room(self, room: str, websocket: WebSocket):
|
|
if room not in self.rooms:
|
|
self.rooms[room] = []
|
|
self.rooms[room].append(websocket)
|
|
|
|
async def remove_from_room(self, room: str, websocket: WebSocket):
|
|
if room in self.rooms:
|
|
self.rooms[room].remove(websocket)
|
|
if len(self.rooms[room]) == 0:
|
|
del self.rooms[room]
|
|
|
|
async def broadcast_to_room(self, room: str, message: str, exclude_websocket: WebSocket = None):
|
|
if room in self.rooms:
|
|
for ws in self.rooms[room]:
|
|
if ws != exclude_websocket:
|
|
try:
|
|
await ws.send_text(message)
|
|
except:
|
|
await self.remove_from_room(room, ws)
|
|
|
|
async def send_to_room(self, room: str, message: str):
|
|
if room in self.rooms:
|
|
for ws in self.rooms[room]:
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
async def send_stream_to_room(self, room: str, message):
|
|
if room in self.rooms:
|
|
for ws in self.rooms[room]:
|
|
try:
|
|
await ws.send_bytes(message)
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
|
|
room_manager = SocketManager()
|