fastapi之WebSockets

WebSockets

WebSockets 是一种在客户端和服务器之间建立持久双向连接的协议。与传统的 HTTP 请求-响应模型不同，WebSocket 允许服务器和客户端可以在连接建立后随时互相发送消息，这使得它非常适合需要实时通信的应用场景，如聊天应用、实时通知、协作编辑等。
FastAPI 内置对 WebSockets 的支持，使得开发实时应用程序变得简单而高效。

基本概念

• **持久连接：**WebSocket 连接一旦建立，客户端和服务器可以在不需要重新建立连接的情况下多次发送和接收数据。
• **双向通信：**服务器可以主动向客户端发送消息，客户端也可以主动向服务器发送消息，形成双向的通信通道。
• **低延迟：**由于不需要为每次通信重新建立连接，WebSocket 通信的延迟较低，非常适合实时应用场景。

FastAPI 中的 WebSocket 支持

FastAPI 提供了对 WebSockets 的原生支持，你可以像定义 HTTP 路由一样定义 WebSocket 路由。

WebSocket 应用示例

示例 1: 简单的 WebSocket 连接

from fastapi import FastAPI, WebSocket

app = FastAPI()

@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: WebSocket: WebSocket 对象用于管理连接、接收和发送数据。
await websocket.accept(): 接受来自客户端的 WebSocket 连接请求。
await websocket.receive_text(): 等待接收来自客户端的文本消息。
await websocket.send_text(): 发送文本消息到客户端。

示例 2: 广播消息的 WebSocket 实现

from fastapi import FastAPI, WebSocket
from typing import List

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)

manager = ConnectionManager()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(f"Client says: {data}")
    except Exception as e:
        manager.disconnect(websocket)

解释

ConnectionManager: 管理所有连接的类，负责接入、断开连接和广播消息。
manager.connect(): 接受新连接并将其添加到活动连接列表中。
manager.disconnect(): 从活动连接列表中移除断开的连接。
manager.broadcast(): 向所有活动连接发送消息。

客户端代码示例

const socket = new WebSocket("ws://localhost:8000/ws");

socket.onmessage = function(event) {
    console.log("Message from server ", event.data);
};

socket.onopen = function(event) {
    socket.send("Hello Server!");
};

完整示例

项目结构

WebSockets_/
│
├── main.py          # 主应用文件
├── static/
│   └── index.html   # 简单的前端客户端
├── requirements.txt # 依赖项
└── README.md        # 项目说明

服务器端代码 (main.py)

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from typing import List
import uvicorn
import os

app = FastAPI()


# Mount the static directory
app.mount("/static", StaticFiles(directory="static"), name="static")


class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


manager = ConnectionManager()


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(f"Client says: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast("A client disconnected")


if __name__ == "__main__":
    uvicorn.run(
        f"{os.path.basename(__file__).split('.')[0]}:app",
        host="127.0.0.1",
        port=8000,
        reload=True,
    )

解释

• ConnectionManager: 负责管理所有的 WebSocket 连接，包括连接、断开连接和消息广播。
• ·/ws 端点·: 当客户端连接到这个端点时，会建立 WebSocket 连接。所有连接的客户端可以通过这个端点发送和接收消息。

简单的前端客户端 (static/index.html)

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>WebSocket Chat</title>
</head>
<body>
    <h1>WebSocket Chat</h1>
    <input id="messageInput" type="text" placeholder="Type a message...">
    <button onclick="sendMessage()">Send</button>
    <ul id="messages"></ul>

    <script>
        const ws = new WebSocket("ws://localhost:8000/ws");
        const messages = document.getElementById('messages');

        ws.onmessage = function(event) {
            const messageItem = document.createElement('li');
            messageItem.textContent = event.data;
            messages.appendChild(messageItem);
        };

        function sendMessage() {
            const input = document.getElementById("messageInput");
            ws.send(input.value);
            input.value = '';
        }
    </script>
</body>
</html>

解释

• 简单的 HTML 页面使用 JavaScript 与 WebSocket 服务器通信。
• 用户在输入框中输入消息并点击发送按钮后，消息会通过 WebSocket 发送到服务器，服务器再将消息广播给所有连接的客户端。

测试

1. 浏览器进入http://127.0.0.1:8000/static/index.html
   然后输入消息 点击send
   
   
2. 关闭网页之后
   

相关代码资源

fastapi之WebSockets

总结

通过 FastAPI 的 WebSocket 支持，开发者可以轻松创建高效的实时应用程序。无论是简单的聊天应用，还是复杂的实时通知系统，FastAPI 提供了一个灵活而强大的平台来满足实时通信的需求。结合前端的 WebSocket API 和 FastAPI 的路由与连接管理功能，开发者可以快速构建现代化的实时交互式应用。