跳转至

FastAPI 处理WebSocket

概要: 利用FastAPI的WebSocket功能主推消息到前端页面

创建时间: 2023.10.08 22:53:27

更新时间: 2023.10.08 23:31:30

安装依赖

Bash
1
2
pip install fastapi
pip install uvicorn

示例代码

文件结构

image.png

./app.py

Python
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import uvicorn
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.websockets import WebSocketDisconnect

app = FastAPI()

# 允许访问静态文件
app.mount("/static", StaticFiles(directory="static"), name="static")

# 用于跟踪所有连接的客户端
connected_clients = []


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    connected_clients.append(websocket)  # 将新连接的客户端添加到列表中
    try:
        while True:
            data = await websocket.receive_text()
            await broadcast_message(data)  # 广播消息给所有连接的客户端
    except WebSocketDisconnect:
        connected_clients.remove(websocket)  # 客户端断开连接时从列表中移除


async def broadcast_message(message: str):
    for client in connected_clients:
        await client.send_text(f"You sent: {message}")


@app.get("/")
async def get():
    with open("./static/index.html") as f:
        return HTMLResponse(content=f.read(), status_code=200)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
  1. 将已建连接的客户端记录在一个列表connected_clients中
  2. 监听endpoint为 /ws 的 WebSocket 事件并进行处理（即调用broadcast_message函数广播收到的消息内容），当客户端断开时，刷新列表connected_clients
  3. 在endpoint为 / 建立HTML服务，供浏览器访问

./static/index.html

HTML
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
<!DOCTYPE html>
<html>
<head>
    <title>WebSocket Example</title>
</head>
<body>
    <h1>WebSocket Example</h1>
    <div id="messages"></div>
    <input type="text" id="input" placeholder="Enter a message" />
    <button onclick="sendMessage()">Send</button>

    <script>
        const socket = new WebSocket('ws://localhost:8000/ws');

        socket.onmessage = function(event) {
            const messages = document.getElementById('messages');
            const message = document.createElement('p');
            message.innerHTML = event.data;
            messages.appendChild(message);
        };

        function sendMessage() {
            const input = document.getElementById('input');
            const message = input.value;
            socket.send(message);
            input.value = '';
        }
    </script>
</body>
</html>
  1. 建立了一个ws://localhost:8000/ws的WebSocket链接
  2. 当WebSocket收到消息时，在HTML中显示
  3. 当浏览器访问此页面，并且用户点击Send按钮时，通过WebSocket发送给后端

使用示例

打开两个浏览器窗口，或者两个浏览器标签（截图示例中上为Chrome，下为Safari）
image.png
接下来，当用户在输入框中输入文字并点击Send按钮发送时，两个浏览器页面将同时刷新收到的消息
FastAPI-ws.gif

参考