FastAPI: WebSocket的用法及举例

引言

在现代网络应用中，实时通信变得越来越重要。WebSocket协议为客户端和服务器之间提供了全双工通信，这为开发实时应用（如聊天应用、在线游戏、股票交易系统等）奠定了基础。FastAPI作为一个高性能的Python框架，提供了对WebSocket的良好支持，使得开发实时应用变得更加简单和高效。

WebSocket简介

WebSocket是一种网络通信协议，允许在单个TCP连接上进行双向交互。它的主要优势包括：

低延迟：WebSocket连接是持久的，避免了频繁重建连接的开销。
全双工通信：服务器可以主动向客户端推送消息，而不仅仅是响应请求。
更少的开销：与HTTP相比，WebSocket在数据传输中的开销更小。

WebSocket工作原理

连接建立：客户端通过HTTP请求发起WebSocket连接，服务器响应并建立连接。
数据传输：一旦连接建立，双方可以随时发送和接收消息。
连接关闭：当会话结束时，可以通过特殊的关闭帧来关闭连接。

FastAPI概述

FastAPI是一个基于Python的现代Web框架，旨在快速构建高性能API。其主要特点包括：

高性能：基于Starlette和Pydantic，FastAPI具有极快的性能。
易于使用：使用Python类型提示，自动生成API文档。
异步支持：内置对异步编程的支持，适合处理高并发任务。

安装FastAPI和相关依赖

要使用FastAPI和WebSocket，我们需要安装FastAPI和一个ASGI服务器（如uvicorn）。可以通过以下命令进行安装：

bashCopy Code
pip install fastapi uvicorn

创建一个简单的WebSocket服务器

下面是创建一个简单WebSocket服务器的步骤：

代码示例

pythonCopy Code
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse

app = FastAPI()

html = """
<!DOCTYPE html>
<html>
    <head>
        <title>WebSocket Test</title>
    </head>
    <body>
        <h1>WebSocket Test</h1>
        <input type="text" id="message" autocomplete="off"/>
        <button onclick="sendMessage()">Send</button>
        <ul id="messages"></ul>
        <script>
            const ws = new WebSocket("ws://localhost:8000/ws");
            ws.onmessage = function(event) {
                const messages = document.getElementById('messages');
                const message = document.createElement('li');
                message.textContent = event.data;
                messages.appendChild(message);
            };
            function sendMessage() {
                const input = document.getElementById("message");
                ws.send(input.value);
                input.value = '';
            }
        </script>
    </body>
</html>
"""

@app.get("/")
async def get():
    return HTMLResponse(html)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(f"Message text was: {data}")

运行服务器

可以使用以下命令启动服务器：

bashCopy Code
uvicorn main:app --reload

在浏览器中访问http://localhost:8000，输入消息并发送，服务器将返回相应的消息。

WebSocket的使用场景

WebSocket适用于多种场景，包括但不限于：

实时聊天应用
在线游戏
股票价格更新
协作编辑工具
实时通知系统

实例：实时聊天应用

功能需求

用户可以加入聊天室
用户可以发送和接收消息
支持多个用户同时在线

代码实现

pythonCopy Code
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from typing import List

app = FastAPI()

clients: List[WebSocket] = []

html = """
<!DOCTYPE html>
<html>
    <head>
        <title>Chat Room</title>
    </head>
    <body>
        <h1>Chat Room</h1>
        <input type="text" id="message" autocomplete="off"/>
        <button onclick="sendMessage()">Send</button>
        <ul id="messages"></ul>
        <script>
            const ws = new WebSocket("ws://localhost:8000/ws");
            ws.onmessage = function(event) {
                const messages = document.getElementById('messages');
                const message = document.createElement('li');
                message.textContent = event.data;
                messages.appendChild(message);
            };
            function sendMessage() {
                const input = document.getElementById("message");
                ws.send(input.value);
                input.value = '';
            }
        </script>
    </body>
</html>
"""

@app.get("/")
async def get():
    return HTMLResponse(html)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    clients.append(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            for client in clients:
                if client != websocket:
                    await client.send_text(data)
    except WebSocketDisconnect:
        clients.remove(websocket)

运行与测试

运行服务器后，打开多个浏览器标签页，输入不同的消息，观察消息是如何在所有连接的客户端之间传递的。

实例：实时数据推送

功能需求

服务器实时推送随机生成的数据到客户端
客户端实时显示数据

代码实现

pythonCopy Code
import random
import asyncio

@app.websocket("/data")
async def data_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = random.randint(1, 100)
        await websocket.send_text(f"Random number: {data}")
        await asyncio.sleep(1)  # 每秒发送一次数据

运行与测试

在浏览器中通过WebSocket连接/data，实时接收随机数字。

错误处理与连接管理

在处理WebSocket时，连接可能会意外断开。以下是如何管理连接状态和处理错误的基本方法。

错误处理示例

pythonCopy Code
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            # 处理接收到的数据
    except WebSocketDisconnect:
        # 处理断开连接
        print("Client disconnected")

安全性考虑

在使用WebSocket时，需要考虑以下安全性问题：

身份验证：确保用户身份有效，可以使用JWT等机制。
跨站请求伪造（CSRF）：实现CORS策略，防止未授权的跨域请求。
数据验证：确保接收到的数据是合法的，避免注入攻击。

总结

FastAPI为WebSocket提供了强大的支持，使得构建实时应用变得简单而高效。无论是聊天应用还是实时数据推送，FastAPI都能很好地满足需求。通过合理的连接管理和安全性考虑，我们可以构建出稳定可靠的实时通信应用。希望本文能够帮助您更好地理解和使用FastAPI中的WebSocket功能。

以上是关于FastAPI和WebSocket用法的详细介绍和实例。如果您有任何问题或反馈，请随时联系我！

FastAPI: WebSocket的用法及举例

目录

引言

WebSocket简介

WebSocket工作原理

FastAPI概述

安装FastAPI和相关依赖

创建一个简单的WebSocket服务器

代码示例

运行服务器

WebSocket的使用场景

实例：实时聊天应用

功能需求

代码实现

运行与测试

实例：实时数据推送

功能需求

代码实现

运行与测试

错误处理与连接管理

错误处理示例

安全性考虑

总结