|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
|
from fastapi.responses import HTMLResponse
|
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
from fastapi.templating import Jinja2Templates
|
|
|
from fastapi.requests import Request
|
|
|
import uvicorn
|
|
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
# 挂载模板目录
|
|
|
templates = Jinja2Templates(directory="templates")
|
|
|
|
|
|
# 活跃连接列表
|
|
|
active_connections = []
|
|
|
|
|
|
# WebSocket 与昵称映射表
|
|
|
usernames = {}
|
|
|
|
|
|
# 首页路由
|
|
|
@app.get("/")
|
|
|
async def get(request: Request):
|
|
|
return templates.TemplateResponse("chat.html", {"request": request})
|
|
|
|
|
|
# WebSocket 路由
|
|
|
@app.websocket("/ws")
|
|
|
async def websocket_endpoint(websocket: WebSocket):
|
|
|
await websocket.accept()
|
|
|
active_connections.append(websocket)
|
|
|
await broadcast_online_count() # 一连上就广播在线人数
|
|
|
|
|
|
try:
|
|
|
# 接收第一个消息作为昵称
|
|
|
username = await websocket.receive_text()
|
|
|
usernames[websocket] = username
|
|
|
|
|
|
await broadcast(f"[系统]{username} 加入了聊天室")
|
|
|
await broadcast_online_count()
|
|
|
|
|
|
while True:
|
|
|
data = await websocket.receive_text()
|
|
|
await broadcast(f"{username}：{data}")
|
|
|
|
|
|
except WebSocketDisconnect:
|
|
|
active_connections.remove(websocket)
|
|
|
left_username = usernames.pop(websocket, "匿名用户")
|
|
|
|
|
|
await broadcast(f"[系统]{left_username} 离开了聊天室")
|
|
|
await broadcast_online_count()
|
|
|
|
|
|
# 广播消息给所有连接
|
|
|
async def broadcast(message: str):
|
|
|
to_remove = []
|
|
|
for connection in active_connections:
|
|
|
try:
|
|
|
await connection.send_text(message)
|
|
|
except:
|
|
|
to_remove.append(connection)
|
|
|
for conn in to_remove:
|
|
|
active_connections.remove(conn)
|
|
|
usernames.pop(conn, None)
|
|
|
|
|
|
# 广播在线人数
|
|
|
async def broadcast_online_count():
|
|
|
message = f"[人数]{len(active_connections)}"
|
|
|
to_remove = []
|
|
|
for conn in active_connections:
|
|
|
try:
|
|
|
await conn.send_text(message)
|
|
|
except:
|
|
|
to_remove.append(conn)
|
|
|
for conn in to_remove:
|
|
|
active_connections.remove(conn)
|
|
|
usernames.pop(conn, None)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
uvicorn.run("chat:app", host="0.0.0.0", port=8202) |
