60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
|
"""
|
|||
|
|
@Time ： 2022/10/12 17:55
|
|||
|
|
@Auth ： 东
|
|||
|
|
@File ：websocket_tool.py
|
|||
|
|
@IDE ：PyCharm
|
|||
|
|
@Motto：ABC(Always Be Coding)
|
|||
|
|
@Desc：
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
import json
|
|||
|
|
from typing import Union, List, Dict
|
|||
|
|
|
|||
|
|
from app.core.common_utils import logger
|
|||
|
|
from app.utils.JSONEncodeTools import MyEncoder
|
|||
|
|
|
|||
|
|
|
|||
|
|
class WebsocketUtil:
|
|||
|
|
def __init__(self):
|
|||
|
|
self.active_connections: List = []
|
|||
|
|
self.active_connections_dist: Dict = {}
|
|||
|
|
|
|||
|
|
def connect(self, ws, id: str):
|
|||
|
|
# 等待连接
|
|||
|
|
msg = ws.receive()
|
|||
|
|
# 存储ws连接对象
|
|||
|
|
self.active_connections.append(ws)
|
|||
|
|
if id in self.active_connections_dist:
|
|||
|
|
self.active_connections_dist[id].append(ws)
|
|||
|
|
else:
|
|||
|
|
ws_list = [ws, ]
|
|||
|
|
self.active_connections_dist[id] = ws_list
|
|||
|
|
|
|||
|
|
def disconnect(self, ws, id):
|
|||
|
|
# ws关闭时 移除ws对象
|
|||
|
|
if ws.closed:
|
|||
|
|
if ws in self.active_connections_dist.values():
|
|||
|
|
self.active_connections.remove(ws)
|
|||
|
|
self.active_connections_dist[id].pop(ws)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
async def send_personal_message(message: str, ws):
|
|||
|
|
# 发送个人消息
|
|||
|
|
await ws.send(message)
|
|||
|
|
|
|||
|
|
def broadcast(self, message: str):
|
|||
|
|
# 广播消息
|
|||
|
|
for connection in self.active_connections:
|
|||
|
|
connection.send(message)
|
|||
|
|
|
|||
|
|
def send_message_proj_json(self, message: Union[str, int, List, Dict], id: str):
|
|||
|
|
# 广播该项目的消息
|
|||
|
|
for connection in self.active_connections_dist[id]:
|
|||
|
|
try:
|
|||
|
|
connection.send(json.dumps(message, cls=MyEncoder, indent=4), )
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("websocket异常断开，{}", e)
|
|||
|
|
self.disconnect(ws=connection, id=id)
|
|||
|
|
|
|||
|
|
|
|||
|
|
manager = WebsocketUtil()
|