FastAPI 学习之路（五十）WebSockets（六）聊天室完善

我们这次只是对于之前的功能做下优化，顺便利用下之前的操作数据的接口，使用下数据库的练习。

在聊天里会有一个上线的概念。上线要通知大家，下线也要通知大家谁离开了，基于此功能我们完善下代码。

首先，我们的登录用户还是用了之前分享中的数据库相关用户，直接调用，之前也是从数据库获取用户，校验密码，密码成功后才让登录。这个在FastApi学习之路（四十七）WebSockets（三）登录后才可以聊天中就已经提到过了。

其次，我们建立了链接管理类，但是我们现在的链接都是没有针对发消息做管理的，我们去扩展下我们之前的发消息的类。

"""
websocket 链接管理
"""from typing import List, Dictfrom starlette.websockets import WebSocketclass ConnectionManager:def __init__(self):"""存放链接"""self.active_connections: List[Dict[str, WebSocket]] = []async def connect(self, user: str, ws: WebSocket):"""链接"""await ws.accept()self.active_connections.append({"user": user, "ws": ws})async def disconnect(self, user: str, ws: WebSocket):"""断开链接，移除"""self.active_connections.remove({"user": user, "ws": ws})@staticmethodasync def send_personal_message(message: str, ws: WebSocket):"""发送所有人消息"""await ws.send_text(message)async def send_other_message(self, message: dict, user: str):"""发送个人消息"""for coon in self.active_connections:if coon["user"] == user:await coon["ws"].send_json(message)async def broadcast(self, data: str):"""广播"""for conn in self.active_connections:await conn["ws"].send_text(data)

增加了广播消息，和增加了发所有消息和针对个人发消息。我们现在实现的是针对所有人和广播。那么我们看下，我们的管理聊天链接的已经完善。

接下来我们就是去实现如何接受消息。

@app.websocket("/items/ws/{user}")
async def websocket_chat_method(websocket: WebSocket,user: str,cookie_or_token: str = Depends(get_cookie_or_token)
):await ws_manager.connect(user, websocket)await ws_manager.broadcast(user + "进入聊天室")try:while True:data = await websocket.receive_text()await ws_manager.send_personal_message(f"你说：{data}", websocket)await ws_manager.broadcast(f"用户{user}说：{data}")except WebSocketDisconnect as e:await ws_manager.disconnect(user, websocket)await ws_manager.broadcast(f"用户{user}离开")

其实很简单，当你登录的时候，就广播一下，某某进入了聊天室，然后就是等消息，那么我们看下前端需要如何改造：

<!DOCTYPE html>
<html>
<head><title>Chat</title>
</head>
<body>
<h1>WebSocket 聊天</h1>
<form action="" onsubmit="sendMessage(event)"><input type="text" id="messageText" autocomplete="off"/><button>Send</button>
</form>
<button onclick="logout()">退出</button>
<ul id='messages'>
</ul>
<script>var  token=window.localStorage.getItem("token")if (token==null ){window.location.href="/login"}var ws = new WebSocket("ws://localhost:8000/items/ws/"+token+"?token="+token);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()}function logout() {window.localStorage.removeItem("token")window.location.href='/login'}
</script></body></html>

其实就是在登录的时候增加了对应人的链接，接下来看下效果：

登录A：