FastAPI WebSocket 聊天室:用戶加入廣播 + 私聊實現
作者:Ss肥魚
在實際聊天室中,我們通常希望有廣播系統消息和私聊功能,而FastAPI + WebSocket 支持這些非常輕松。
在實際聊天室中,我們通常希望有以下功能:
- 廣播系統消息:某個用戶加入/離開聊天室,通知所有人。
- 私聊功能:用戶可以給指定用戶發送私密消息。
FastAPI + WebSocket 支持這些非常輕松。
目錄結構
fastapi-chat/
├── main.py
└── chat_manager.py
依賴安裝
pip install fastapi uvicorn
(1) chat_manager.py:升級連接管理器,支持用戶身份
# chat_manager.py
from fastapi import WebSocket
from typing import Dict
class ConnectionManager:
def __init__(self):
self.active_connections: Dict[str, WebSocket] = {}
async def connect(self, username: str, websocket: WebSocket):
await websocket.accept()
self.active_connections[username] = websocket
await self.broadcast(f"?? 用戶【{username}】加入了聊天室")
def disconnect(self, username: str):
if username in self.active_connections:
del self.active_connections[username]
async def send_personal_message(self, message: str, to_user: str):
if to_user in self.active_connections:
await self.active_connections[to_user].send_text(message)
async def broadcast(self, message: str):
for ws in self.active_connections.values():
await ws.send_text(message)
(2) main.py:WebSocket 支持用戶名、私聊
# main.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
from chat_manager import ConnectionManager
app = FastAPI()
manager = ConnectionManager()
@app.websocket("/ws/chat/")
async def websocket_endpoint(websocket: WebSocket, username: str = Query(...)):
await manager.connect(username, websocket)
try:
while True:
data = await websocket.receive_json()
msg_type = data.get("type", "broadcast")
content = data.get("message")
to_user = data.get("to")
if msg_type == "private" and to_user:
await manager.send_personal_message(f"?? 私信【{username}】→【{to_user}】:{content}", to_user)
else:
await manager.broadcast(f"?? 【{username}】:{content}")
except WebSocketDisconnect:
manager.disconnect(username)
await manager.broadcast(f"? 用戶【{username}】離開了聊天室")
HTML 前端示例(支持私聊發送)
<!-- index.html -->
<!DOCTYPE html>
<html>
<head><title>聊天</title></head>
<body>
<h2>FastAPI 聊天室</h2>
<input id="username" placeholder="用戶名" />
<button onclick="connect()">連接</button>
<div>
<input id="to" placeholder="發給誰(留空為群發)" />
<input id="message" placeholder="消息內容" />
<button onclick="send()">發送</button>
</div>
<ul id="messages"></ul>
<script>
let ws = null;
function connect() {
const username = document.getElementById("username").value;
ws = new WebSocket("ws://localhost:8000/ws/chat/?username=" + username);
ws.onmessage = event => {
const li = document.createElement("li");
li.innerText = event.data;
document.getElementById("messages").appendChild(li);
};
}
function send() {
const msg = document.getElementById("message").value;
const to = document.getElementById("to").value;
const type = to ? "private" : "broadcast";
ws.send(JSON.stringify({ type, message: msg, to }));
}
</script>
</body>
</html>
功能回顧
- 新用戶連接 → 自動廣播 ?? 用戶加入
- 用戶斷開連接 → 自動廣播 ? 用戶離開
- 留空接收人 → 群發廣播消息
- 指定 to 字段 → 發送私聊
責任編輯:趙寧寧
來源:
Ssoul肥魚