可以在 FastAPI 中使用 WebSockets。
安装 websockets
把 websockets 添加到项目中。它是一个让使用 WebSocket 协议更容易的 Python 库:
$ uv add websockets
---> 100%
WebSockets 客户端
生产环境
生产系统的前端可能使用 React、Vue.js 或 Angular 等现代框架。要通过 WebSockets 与后端通信,通常使用前端框架的工具。也可能有原生移动应用,直接用原生代码连接 WebSocket 后端;或者以其他方式与 WebSocket 端点通信。
本示例使用一个非常简单的 HTML 文档,其中包含 JavaScript,全部放在一个长字符串中。这不是最优做法,不应在生产环境中这样组织页面;生产环境应选择前述方式之一。但为了专注展示服务器端 WebSockets,并提供可用示例,这是最简单的方式:
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
app = FastAPI()
html = """
<!DOCTYPE html>
<html>
<head>
<title>Chat</title>
</head>
<body>
<h1>WebSocket Chat</h1>
<form action="" onsubmit="sendMessage(event)">
<input type="text" id="messageText" autocomplete="off"/>
<button>Send</button>
</form>
<ul id='messages'>
</ul>
<script>
var ws = new WebSocket("ws://localhost:8000/ws");
ws.onmessage = function(event) {
var messages = document.getElementById('messages')
var message = document.createElement('li')
var content = document.createTextNode(event.data)
message.appendChild(content)
messages.appendChild(message)
};
function sendMessage(event) {
var input = document.getElementById("messageText")
ws.send(input.value)
input.value = ''
event.preventDefault()
}
</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}")
创建 websocket
在 FastAPI 应用中创建 websocket。上例使用 @app.websocket("/ws"),端点接收 websocket: WebSocket。
等待并发送消息
在 WebSocket 路由中,可以用 await 等待消息并发送消息。上例先调用 await websocket.accept() 接受连接,再在循环中用 await websocket.receive_text() 接收文本、用 await websocket.send_text(...) 回复。除了文本,还可以接收和发送二进制与 JSON 数据。
试一试
把代码放入 main.py,然后运行应用:
$ uv run fastapi dev
<span style="color: green;">INFO</span>: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
在浏览器打开 http://127.0.0.1:8000。页面如下:

在输入框输入并发送消息:

FastAPI 应用通过 WebSockets 回复:

可以连续发送和接收多条消息:

所有消息都使用同一个 WebSocket 连接。
使用 Depends 和其他依赖项
在 WebSocket 端点中,可以从 fastapi 导入并使用 Depends、Security、Cookie、Header、Path 和 Query。它们与其他 FastAPI 端点/路径操作中的工作方式相同:
from typing import Annotated
from fastapi import (
Cookie,
Depends,
FastAPI,
Query,
WebSocket,
WebSocketException,
status,
)
from fastapi.responses import HTMLResponse
app = FastAPI()
html = """
<!DOCTYPE html>
<html>
<head>
<title>Chat</title>
</head>
<body>
<h1>WebSocket Chat</h1>
<form action="" onsubmit="sendMessage(event)">
<label>Item ID: <input type="text" id="itemId" autocomplete="off" value="foo"/></label>
<label>Token: <input type="text" id="token" autocomplete="off" value="some-key-token"/></label>
<button onclick="connect(event)">Connect</button>
<hr>
<label>Message: <input type="text" id="messageText" autocomplete="off"/></label>
<button>Send</button>
</form>
<ul id='messages'>
</ul>
<script>
var ws = null;
function connect(event) {
var itemId = document.getElementById("itemId")
var token = document.getElementById("token")
ws = new WebSocket("ws://localhost:8000/items/" + itemId.value + "/ws?token=" + token.value);
ws.onmessage = function(event) {
var messages = document.getElementById('messages')
var message = document.createElement('li')
var content = document.createTextNode(event.data)
message.appendChild(content)
messages.appendChild(message)
};
event.preventDefault()
}
function sendMessage(event) {
var input = document.getElementById("messageText")
ws.send(input.value)
input.value = ''
event.preventDefault()
}
</script>
</body>
</html>
"""
@app.get("/")
async def get():
return HTMLResponse(html)
async def get_cookie_or_token(
websocket: WebSocket,
session: Annotated[str | None, Cookie()] = None,
token: Annotated[str | None, Query()] = None,
):
if session is None and token is None:
raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
return session or token
@app.websocket("/items/{item_id}/ws")
async def websocket_endpoint(
*,
websocket: WebSocket,
item_id: str,
q: int | None = None,
cookie_or_token: Annotated[str, Depends(get_cookie_or_token)],
):
await websocket.accept()
while True:
data = await websocket.receive_text()
await websocket.send_text(
f"Session cookie or query token value is: {cookie_or_token}"
)
if q is not None:
await websocket.send_text(f"Query parameter q is: {q}")
await websocket.send_text(f"Message text was: {data}, for item ID: {item_id}")
尝试带依赖项的 WebSockets
运行应用:
$ uv run fastapi dev
<span style="color: green;">INFO</span>: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
在浏览器打开 http://127.0.0.1:8000。可以设置用于路径的“Item ID”,以及作为查询参数的“Token”。
接着连接 WebSocket,然后收发消息:

处理断开连接和多个客户端
WebSocket 连接关闭时,await websocket.receive_text() 会引发 WebSocketDisconnect 异常,可以捕获并处理:
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
app = FastAPI()
html = """
<!DOCTYPE html>
<html>
<head>
<title>Chat</title>
</head>
<body>
<h1>WebSocket Chat</h1>
<h2>Your ID: <span id="ws-id"></span></h2>
<form action="" onsubmit="sendMessage(event)">
<input type="text" id="messageText" autocomplete="off"/>
<button>Send</button>
</form>
<ul id='messages'>
</ul>
<script>
var client_id = Date.now()
document.querySelector("#ws-id").textContent = client_id;
var ws = new WebSocket(`ws://localhost:8000/ws/${client_id}`);
ws.onmessage = function(event) {
var messages = document.getElementById('messages')
var message = document.createElement('li')
var content = document.createTextNode(event.data)
message.appendChild(content)
messages.appendChild(message)
};
function sendMessage(event) {
var input = document.getElementById("messageText")
ws.send(input.value)
input.value = ''
event.preventDefault()
}
</script>
</body>
</html>
"""
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.get("/")
async def get():
return HTMLResponse(html)
@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: int):
await manager.connect(websocket)
try:
while True:
data = await websocket.receive_text()
await manager.send_personal_message(f"You wrote: {data}", websocket)
await manager.broadcast(f"Client #{client_id} says: {data}")
except WebSocketDisconnect:
manager.disconnect(websocket)
await manager.broadcast(f"Client #{client_id} left the chat")
尝试在多个浏览器标签页打开应用,从这些标签页发送消息,然后关闭其中一个。关闭会触发 WebSocketDisconnect,所有其他客户端会收到类似消息:
Client #1596980209979 left the chat
更多信息
更多选项见 Starlette 文档:WebSocket 类与基于类的 WebSocket 处理。











暂无评论内容