105 lines
2.7 KiB
Python
105 lines
2.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
|
||
import os
|
||
|
||
from app.core.common_utils import logger
|
||
from app.utils.JSONEncodeTools import MyEncoder
|
||
from configs.global_var import *
|
||
from app.utils.redis_config import redis_client
|
||
from multiprocessing import Process, Queue
|
||
import pickle
|
||
|
||
|
||
class WebsocketUtil:
|
||
|
||
def __init__(self):
|
||
# self.active_connections = multiprocessing.Manager().list()
|
||
self.active_connections_dist = []
|
||
|
||
|
||
def connect(self, ws, id: str):
|
||
# 等待连接
|
||
msg = ws.receive()
|
||
# 存储ws连接对象
|
||
if os.path.exists(f"{id}.pkl"):
|
||
ws_list = read(id=id)
|
||
ws_list.append(ws)
|
||
data = WsFile(id=id, ws=ws_list)
|
||
write(id=id, ws=data)
|
||
else:
|
||
ws_list = [ws, ]
|
||
data = WsFile(id=id, ws=ws_list)
|
||
write(id=id, ws=data)
|
||
print("--;;-----------", ws_list)
|
||
|
||
|
||
|
||
def disconnect(self, ws, id):
|
||
# 删除连接
|
||
if os.path.exists(f"{id}.pkl"):
|
||
os.remove(f"{id}.pkl")
|
||
|
||
@staticmethod
|
||
async def send_personal_message(message: str, ws):
|
||
# 发送个人消息
|
||
await ws.send(message)
|
||
|
||
# def broadcast(self, message: str):
|
||
# # 广播消息
|
||
# # global active_connections
|
||
# active_connections = redis_client.get_redis().get("active_connections")
|
||
# if active_connections is not None:
|
||
# active_connections = json.loads(active_connections)
|
||
# for connection in active_connections:
|
||
# connection.send(message)
|
||
|
||
def send_message_proj_json(self, message: Union[str, int, List, Dict], id: str):
|
||
# 广播该项目的消息
|
||
active_connections = read(id=id)
|
||
print("===================", type(active_connections), active_connections[0])
|
||
for connection in active_connections:
|
||
try:
|
||
connection.send(json.dumps(message, cls=MyEncoder, indent=4), )
|
||
except Exception as e:
|
||
logger.error("websocket异常断开,{}", e)
|
||
|
||
|
||
manager = WebsocketUtil()
|
||
|
||
|
||
def write(id: str, ws: List):
|
||
print(f"序列化对象{ws}")
|
||
with open(f"{id}.pkl", "wb") as f:
|
||
pickle.dump(ws, f)
|
||
|
||
|
||
def read(id: str):
|
||
|
||
with open(f"{id}.pkl", "rb") as f:
|
||
wss = pickle.load(f)
|
||
print(wss)
|
||
print(f"反序列化对象{wss}")
|
||
return wss
|
||
|
||
|
||
class WsFile:
|
||
def __init__(self, id: str, ws: List) -> None:
|
||
self.id = id
|
||
self.ws = ws
|
||
|
||
def __getstate__(self):
|
||
pickled = {"id": id}
|
||
return pickled
|
||
|
||
def __setstate(self, pickled_dict):
|
||
self.__init__(pickled_dict['id'])
|