10.2 WebSocket 端点定义

FastAPI WebSocket 端点定义入门教程 | Python 初学者指南

FastAPI 教程

本教程面向Python初学者，详细讲解如何在FastAPI中定义WebSocket端点，无需Web开发经验。通过生活化类比和简单示例，帮助您快速上手，建立信心。

FastAPI WebSocket 端点定义入门教程

欢迎来到FastAPI的WebSocket世界！如果您是Python新手，但对构建实时应用感兴趣，这个教程就是为您量身定做的。我们将跳过复杂的理论，直接动手实践，让您快速体验到成功的喜悦。

什么是WebSocket？生活化类比

想象一下，HTTP请求就像发短信：您发送一条消息，然后等待回复。而WebSocket更像是打视频电话：一旦连接建立，双方可以实时、双向地交流，不断发送数据。在Web开发中，WebSocket常用于聊天应用、实时游戏或股票行情更新等场景。

FastAPI是一个现代、快速的Python Web框架，它让定义WebSocket端点变得超级简单！

准备工作：安装所需工具

在开始之前，请确保您已经安装了Python（建议3.7及以上版本）。然后，打开终端或命令行，运行以下命令安装FastAPI和服务器：

pip install fastapi uvicorn websockets

fastapi: 核心框架
uvicorn: 一个轻量级的ASGI服务器，用于运行FastAPI应用
websockets: WebSocket库，FastAPI依赖它来处理WebSocket连接

安装完成后，您就准备好了一切！

第一步：创建您的第一个FastAPI应用

让我们从一个最简单的FastAPI应用开始，这样您可以先看到框架的基本结构。创建一个新文件，比如 main.py，并输入以下代码：

from fastapi import FastAPI

app = FastAPI()

@app.get("/")
def home():
    return {"message": "欢迎来到FastAPI世界！"}

运行这个应用：

uvicorn main:app --reload

在浏览器中打开 http://127.0.0.1:8000，您应该看到 {"message": "欢迎来到FastAPI世界！"}。恭喜！您已经成功运行了一个FastAPI应用。这是一个HTTP端点，但现在我们切换到WebSocket。

第二步：定义您的第一个WebSocket端点

现在，让我们添加一个WebSocket端点。WebSocket端点使用 @app.websocket 装饰器来定义，类似于HTTP的 @app.get。在 main.py 文件中添加以下代码：

from fastapi import FastAPI, WebSocket

app = FastAPI()

# 已有的HTTP端点
@app.get("/")
def home():
    return {"message": "欢迎来到FastAPI世界！"}

# 新增的WebSocket端点
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()  # 接受WebSocket连接
    await websocket.send_text("WebSocket连接已建立！请输入您的消息。")  # 发送欢迎消息
    while True:
        data = await websocket.receive_text()  # 接收客户端发送的文本消息
        # 简单处理：将消息原样返回，就像回声一样
        await websocket.send_text(f"您说：{data}")

代码解析：成功时刻到来！

导入模块：我们导入了 WebSocket 类，这是FastAPI中处理WebSocket连接的关键。
定义端点：@app.websocket("/ws") 告诉FastAPI，当客户端连接到 ws://127.0.0.1:8000/ws 时，调用 websocket_endpoint 函数。路径 /ws 是任意的，您可以根据需要更改。
异步函数：WebSocket操作通常是异步的，所以我们用 async def。如果您不熟悉异步编程，别担心——FastAPI会自动处理大部分细节，您只需使用 await 关键字。
连接流程：
- await websocket.accept(): 这是建立连接的步骤，就像接起电话。
- await websocket.send_text(...): 发送一条欢迎消息给客户端。
- while True: 无限循环，持续监听客户端消息。
- data = await websocket.receive_text(): 接收客户端发来的文本消息。
- await websocket.send_text(...): 将消息原样发送回去，实现回声效果。

运行应用：

uvicorn main:app --reload

现在，我们需要测试WebSocket连接。您可以使用浏览器开发者工具或一个简单的前端页面。为了方便，这里提供一个HTML文件来测试。创建 test_websocket.html 文件：

<!DOCTYPE html>
<html>
<head>
    <title>测试WebSocket</title>
</head>
<body>
    <input type="text" id="messageInput" placeholder="输入消息">
    <button onclick="sendMessage()">发送</button>
    <div id="response"></div>

    <script>
        const socket = new WebSocket('ws://127.0.0.1:8000/ws');
        socket.onopen = function(event) {
            console.log('WebSocket连接已打开');
        };
        socket.onmessage = function(event) {
            document.getElementById('response').innerHTML += event.data + '<br>';
        };
        function sendMessage() {
            const input = document.getElementById('messageInput');
            socket.send(input.value);
            input.value = '';
        }
    </script>
</body>
</html>

用浏览器打开这个HTML文件，输入消息并点击发送，您会看到服务器回显的消息。看到实时交互了吗？这就是WebSocket的魅力！

第三步：扩展示例——简单聊天室

为了建立更多信心，让我们扩展一下，模拟一个简单聊天室。修改 main.py 的WebSocket部分：

from fastapi import FastAPI, WebSocket
from typing import List

app = FastAPI()

# 存储所有连接的WebSocket客户端
connections: List[WebSocket] = []

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    connections.append(websocket)  # 将新连接添加到列表中
    try:
        while True:
            data = await websocket.receive_text()
            # 广播消息给所有连接的客户端
            for connection in connections:
                await connection.send_text(f"用户说：{data}")
    except Exception as e:
        connections.remove(websocket)  # 如果连接断开，移除它
        print(f"连接断开: {e}")

这个例子中，我们维护了一个客户端列表，当一个用户发送消息时，所有其他用户都会收到。这就是实时聊天的基础！

运行应用并测试多个浏览器标签页打开 test_websocket.html，您会看到消息在所有标签页间同步。

总结与下一步

恭喜！您已经成功学会了在FastAPI中定义WebSocket端点。关键点回顾：

WebSocket类比：像实时电话，双向通信。
定义端点：使用 @app.websocket 装饰器，路径自定义。
异步操作：用 async/await 处理连接和数据交换。

现在，您可以尝试更复杂的应用，比如添加用户身份验证或处理JSON数据。FastAPI文档有更多高级示例，但先从这里开始，建立信心最重要。

动手任务：修改代码，让回声服务器只返回消息的大写版本。运行并测试一下！

记住，实践是最好的老师。继续探索，享受编程的乐趣！

本教程使用FastAPI 0.104+版本，确保您的环境兼容。如有问题，请检查Python版本和包安装。

上一章 10.1 WebSocket 协议基础

下一章 10.3 双向实时通信实现