RODY/app/utils/websocket_tool.py

"""
@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)
        else:
            ws_list = [ws, ]
        write(id=id, 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("===================", len(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):
    with open(f"{id}.pkl", "wb") as f:
        wsdump = pickle.dump(ws, f)


def read(id: str):
    with open(f"{id}.pkl", "rb") as f:
        wss = pickle.load(f)
        print(wss)
    return wss