31 lines
1.0 KiB
Python
31 lines
1.0 KiB
Python
|
from fastapi import WebSocket
|
||
|
|
from fastapi import FastAPI
|
||
|
|
from starlette.websockets import WebSocketState
|
||
|
|
|
||
|
|
from utils.websocket_server import room_manager
|
||
|
|
|
||
|
|
|
||
|
|
def websocket_config(app: FastAPI):
|
||
|
|
|
||
|
|
@app.websocket("/{room}")
|
||
|
|
async def websocket_room(websocket: WebSocket, room: str):
|
||
|
|
"""
|
||
|
|
websocket 房间管理
|
||
|
|
:param websocket:
|
||
|
|
:param room:
|
||
|
|
:return:
|
||
|
|
"""
|
||
|
|
await websocket.accept()
|
||
|
|
await room_manager.add_to_room(room, websocket)
|
||
|
|
try:
|
||
|
|
while True:
|
||
|
|
data = await websocket.receive_text()
|
||
|
|
await room_manager.broadcast_to_room(room, data, exclude_websocket=websocket)
|
||
|
|
except Exception as e:
|
||
|
|
if websocket.client_state != WebSocketState.DISCONNECTED:
|
||
|
|
await websocket.close(code=1000)
|
||
|
|
finally:
|
||
|
|
await room_manager.remove_from_room(room, websocket)
|
||
|
|
if websocket.client_state != WebSocketState.DISCONNECTED:
|
||
|
|
await websocket.close(code=1001)
|