代码收藏家技术教程 2024-10-20

Python 中的 WebSocket 使用

什么是 WebSocket？

WebSocket 是一种网络通信协议，能够在客户端和服务器之间建立全双工（双向）通信通道。这种连接使得双方可以在无需重新建立连接的情况下随时发送数据，适合需要实时数据传输的场景。

WebSocket 的原理

握手阶段：WebSocket 通信始于客户端发送一个特殊的 HTTP 请求，用以请求服务器升级协议为 WebSocket。如果服务器同意，会返回一个状态码 101，表示协议升级成功，接下来的通信会使用 WebSocket 协议。
数据传输阶段：在连接建立后，客户端和服务器之间可以自由地进行双向数据传输。数据被打包成帧（frame）后，通过底层的 TCP 连接传输。
关闭连接：当通信完成或出现异常时，任意一方可以发送一个关闭帧来终止连接。

WebSocket 使用场景

实时聊天应用：如在线聊天室、即时通讯工具。

实时数据更新：如股票价格、体育比赛直播、在线游戏状态。

多人协作：如实时协作编辑工具、在线白板。

实时通知：如在线通知系统、推送服务。

Python 中的 WebSocket 库和 API

在 Python 中，常用的 WebSocket 库包括 websockets 和 FastAPI 等。

1. `websockets` 库

websockets 是一个用于构建 WebSocket 服务器和客户端的库。它提供了简单的 API 来建立 WebSocket 连接和发送接收消息。

安装：

pip install websockets

API 示例：

服务器端代码：

import asyncio
import websockets

# WebSocket 服务器处理函数
async def echo(websocket, path):
    async for message in websocket:
        await websocket.send(f"Echo: {message}")

# 启动 WebSocket 服务器
start_server = websockets.serve(echo, "localhost", 8765)

# 运行事件循环
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()

websockets.serve(handler, host, port): 创建 WebSocket 服务器。

websocket.send(data): 发送消息到客户端。

websocket.recv(): 接收客户端消息。

客户端代码：

import asyncio
import websockets

# 客户端处理函数
async def hello():
    async with websockets.connect("ws://localhost:8765") as websocket:
        await websocket.send("Hello, WebSocket!")
        response = await websocket.recv()
        print(f"Received: {response}")

# 运行客户端
asyncio.get_event_loop().run_until_complete(hello())

websockets.connect(uri): 连接到 WebSocket 服务器。

websocket.send(data): 发送消息到服务器。

websocket.recv(): 接收服务器消息。

2. `FastAPI` 库中的 WebSocket

FastAPI 是一个现代的、快速的 Web 框架，支持异步操作，并且原生支持 WebSocket。

安装：

pip install fastapi uvicorn

API 示例：

服务器端代码：

from fastapi import FastAPI, WebSocket

app = FastAPI()

# WebSocket 端点
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()  # 接受 WebSocket 连接
    while True:
        data = await websocket.receive_text()  # 接收消息
        await websocket.send_text(f"Message text was: {data}")  # 发送消息

app.websocket("/ws"): 定义一个 WebSocket 端点，路径为 /ws。

websocket.accept(): 接受客户端的 WebSocket 连接。

websocket.receive_text(): 接收来自客户端的文本消息。

websocket.send_text(data): 发送文本消息给客户端。

运行服务器：

uvicorn myapp:app --reload

在上面的命令中，myapp 是你的 Python 文件的名字，app 是 FastAPI 应用实例的名字。

场景应用示例：实时聊天应用

场景描述：用户可以在同一个聊天室中互相发送消息，消息会实时显示在所有用户的屏幕上。

实现思路：

使用 WebSocket 进行实时双向通信。
将所有连接的客户端保存在一个列表中，方便向所有客户端广播消息。
当一个客户端发送消息时，服务器将消息转发给所有其他客户端。

代码示例：

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

clients = []  # 存储所有连接的客户端

@app.websocket("/chat")
async def chat(websocket: WebSocket):
    await websocket.accept()
    clients.append(websocket)  # 将新连接的客户端添加到列表中
    try:
        while True:
            data = await websocket.receive_text()
            for client in clients:
                await client.send_text(f"User says: {data}")
    except WebSocketDisconnect:
        clients.remove(websocket)  # 移除断开连接的客户端

运行方式：
通过 uvicorn 运行服务器，并使用浏览器或其他 WebSocket 客户端工具连接到 /chat 端点。每个连接的用户发送的消息都会实时广播给所有其他用户。