什么是 WebSocket？

WebSocket 是一种在单个 TCP 连接上进行全双工通信的协议，它允许服务器主动向客户端推送数据，非常适合需要实时通信的应用，

• 聊天应用
• 实时数据仪表盘
• 在线游戏
• 协作编辑工具

使用 websockets 库 (推荐)

websockets 是 Python 中最流行、功能最全面的 WebSocket 客户端和服务器库之一，它提供了异步 API，性能优异。

安装

你需要安装这个库：

pip install websockets

客户端示例

这是一个简单的 WebSocket 客户端，它会连接到一个公共的测试服务器（wss://echo.websocket.org），发送一条消息，然后打印收到的回复。

import asyncio
import websockets
# 目标 WebSocket 服务器地址
# 这是一个公共的回声服务器，它会把你发送的内容原样返回
URI = "wss://echo.websocket.org"
async def hello():
    # 使用 async with 语句建立连接
    async with websockets.connect(URI) as websocket:
        # 发送一条消息到服务器
        name = "Python WebSocket 客户端"
        await websocket.send(f"Hello from {name}!")
        print(f"已发送: Hello from {name}!")
        # 等待并接收服务器的回复
        # 使用 await 等待直到收到消息
        greeting = await websocket.recv()
        print(f"已收到: {greeting}")
# 运行异步主函数
asyncio.run(hello())

代码解释：

1. import asyncio: 导入异步编程的核心库。
2. import websockets: 导入 websockets 库。
3. URI: 定义要连接的 WebSocket 服务器的地址。wss:// 表示加密的 WebSocket 连接（类似于 HTTPS）。
4. async def hello(): 定义一个异步函数，所有网络操作都必须在异步函数中进行。
5. async with websockets.connect(URI) as websocket:: 这是建立连接的核心。async with 会自动处理连接的建立和关闭。websocket 对象代表了这个连接。
6. await websocket.send(...): 异步发送数据。
7. await websocket.recv(): 异步接收数据，程序会在这里暂停，直到收到服务器的消息。
8. asyncio.run(hello()): 启动并运行异步事件循环，执行 hello() 函数。

使用 websocket-client 库

如果你需要一个更接近传统同步编程体验的库,websocket-client 是一个不错的选择，它同时支持同步和异步模式。

安装

pip install websocket-client

客户端示例 (同步模式)

import websocket
import threading
import time
# 目标 WebSocket 服务器地址
URI = "wss://echo.websocket.org"
def on_message(ws, message):
    """当收到消息时调用此函数"""
    print(f"已收到: {message}")
def on_error(ws, error):
    """当发生错误时调用此函数"""
    print(f"发生错误: {error}")
def on_close(ws, close_status_code, close_msg):
    """当连接关闭时调用此函数"""
    print("### 连接已关闭 ###")
def on_open(ws):
    """当连接成功建立时调用此函数"""
    print("连接已建立")
    # 在新线程中发送消息，避免阻塞主线程
    def run(*args):
        for i in range(3):
            time.sleep(1)
            ws.send(f"这是第 {i+1} 条消息")
            print(f"已发送: 这是第 {i+1} 条消息")
        time.sleep(1)
        ws.close() # 关闭连接
        print("线程结束")
    threading.Thread(target=run).start()
if __name__ == "__main__":
    # 创建 WebSocket 对象
    ws = websocket.WebSocketApp(
        URI,
        on_open=on_open,
        on_message=on_message,
        on_error=on_error,
        on_close=on_close
    )
    # 运行 WebSocket，这会阻塞直到连接关闭
    # setinterval=0.1 表示每 0.1 秒检查一次新消息
    ws.run_forever(setinterval=0.1)

代码解释：

1. 这个库通过回调函数来处理事件：on_open, on_message, on_error, on_close。
2. WebSocketApp 创建了一个应用实例，并将回调函数绑定到相应的事件上。
3. ws.run_forever() 启动客户端，它会一直运行，直到连接关闭，这是一个阻塞调用。

使用 aiohttp 库

aiohttp 是一个功能强大的异步 HTTP 客户端/服务器框架，它也内置了对 WebSocket 的支持，如果你已经在使用 aiohttp 处理 HTTP 请求，那么使用它来处理 WebSocket 会非常方便。

安装

pip install aiohttp

客户端示例

import asyncio
import aiohttp
# 目标 WebSocket 服务器地址
URI = "wss://echo.websocket.org"
async def fetch_aiohttp():
    # 创建一个 ClientSession，这是 aiohttp 进行所有操作的基础
    async with aiohttp.ClientSession() as session:
        # 使用 session.ws_connect() 建立 WebSocket 连接
        async with session.ws_connect(URI) as ws:
            # 发送一条消息
            await ws.send_str("Hello from aiohttp!")
            print("已发送: Hello from aiohttp!")
            # 接收一条消息
            # ws.receive() 会返回一个消息对象
            # 我们使用 text=True 来解码为文本
            response = await ws.receive(text=True)
            print(f"已收到: {response.data}")
# 运行异步主函数
asyncio.run(fetch_aiohttp())

代码解释：

1. aiohttp.ClientSession() 是管理连接池和会话上下文的标准方式。
2. session.ws_connect() 用于建立 WebSocket 连接。
3. ws.send_str() 用于发送文本消息，还有 send_json()、send_bytes() 等方法。
4. await ws.receive(text=True) 接收消息。response.data 包含消息内容。

总结与对比

如何选择？

• 如果你是初学者或追求简洁高效：首选 websockets，它的 API 非常直观，是社区的事实标准。
• 如果你喜欢传统的同步编程风格：使用 websocket-client 的同步模式。
• 如果你的项目已经在使用 aiohttp：那么直接使用 aiohttp 的 WebSocket 功能，保持技术栈统一。

希望这个详细的指南能帮助你开始在 Python 中使用 WebSocket！