Python websocket服务/客户端实现示例

0.安装依赖

pip install websockets

1.服务端

import asyncio
from typing import AsyncIterable, Dict, Iterable, Set, Union

import websockets
import websockets.protocol
from loguru import logger
from websockets.typing import Data


class WebSocketServer:

    def __init__(self):
        self.conn_set: Set[websockets.WebSocketServerProtocol] = set()
        self.conn_map: Dict[websockets.WebSocketServerProtocol, Dict[str, int]] = dict()

    def register_conn(self, ws_conn: websockets.WebSocketServerProtocol):
        """
        注册连接
        """
        conn_id = id(ws_conn)
        if ws_conn in self.conn_set:
            logger.warning(f"连接已存在:{conn_id}")
            return
        self.conn_set.add(ws_conn)
        self.conn_map[ws_conn] = {"id": conn_id}
        logger.info(f"成功注册连接:{conn_id}")

    def unregister_conn(self, ws_conn: websockets.WebSocketServerProtocol):
        """
        注销链接
        """
        if ws_conn not in self.conn_set:
            logger.warning(f"连接不存在,无法注销")
            return
        self.conn_set.remove(ws_conn)
        conn_id = self.conn_map[ws_conn]["id"]
        del self.conn_map[ws_conn]
        logger.info(f"成功注销连接:{conn_id}")

    async def broadcast(
        self,
        content: Union[Data, Iterable[Data], AsyncIterable[Data]],
    ):
        """
        content: 类型是字符串或字节,或是输出它们的可迭代对象
        广播消息
        """
        if not self.conn_set:
            return
        # 给所有连接发数据,并等待
        # send方法会自动判断数据类型,并带上不同的数据帧,让给客户端做辨别
        # 源码中分别是:
        #   frames.Opcode.TEXT和frames.Opcode.BINARY
        await asyncio.wait([ws_conn.send(content) for ws_conn in self.conn_set])

    async def close_conn(
        self,
        ws_conn: websockets.WebSocketServerProtocol,
        code: int,
        msg: str,
    ):
        """
        关闭连接,并尝试向服务端发送关闭消息
        """
        conn_id = id(ws_conn)
        if ws_conn.state != websockets.protocol.OPEN:
            logger.warning(f"要求关闭的连接{conn_id}不是OPEN状态")
            return
        await ws_conn.close(code, msg)

    async def handler(self, ws_conn: websockets.WebSocketServerProtocol, path: str):
        """
        处理单个连接
        """
        self.register_conn(ws_conn)
        try:
            async for content in ws_conn:
                resp_content = f"服务端收到消息:{content}"
                await ws_conn.send(resp_content)
                ...
        except websockets.exceptions.ConnectionClosedError:
            ...
        except Exception:
            ...
        finally:
            self.unregister_conn(ws_conn)


def start_server(server: WebSocketServer, host: str, port: int):
    """
    启动服务端
    """

    async def main():
        async with websockets.serve(server.handler, host, port):
            logger.info("ws服务启动")
            await asyncio.Future()

    asyncio.run(main())


if __name__ == "__main__":
    start_server(server=WebSocketServer(), host="localhost", port=10002)

2.客户端

from typing import AsyncIterable, Iterable, Optional, Union

import websockets
from loguru import logger
from websockets import WebSocketClientProtocol
from websockets.typing import Data


class WebSocketClient:

    def __init__(self, uri: str):
        self.uri: str = uri
        self.ws_conn: Optional[WebSocketClientProtocol] = None

    async def connect(self) -> WebSocketClientProtocol:
        """
        连接到 WebSocket 服务器
        """
        self.ws_conn = await websockets.connect(self.uri)
        logger.info("连接到websocket服务器")
        return self.ws_conn

    async def send_data(self, data: Union[Data, Iterable[Data], AsyncIterable[Data]]):
        """
        发送数据
        """
        assert self.ws_conn, "未建立连接"
        await self.ws_conn.send(data)

    async def recv_data(self):
        """
        接收数据
        """
        assert self.ws_conn, "未建立连接"
        try:
            async for data in self.ws_conn:
                ...
        except websockets.exceptions.ConnectionClosed:
            ...

    async def close(self):
        """
        关闭连接
        """
        assert self.ws_conn, "未建立连接"
        await self.ws_conn.close()