本次是WebSockets的第六次分享了，我们这次只是对于之前的功能做下优化，顺便利用下之前的操作数据的接口，使用下数据库的练习。
我们都知道，在聊天室里面会有一个上线的概念。上线了要通知大家。然后下线了要告诉大家谁离开了。发小小大家要都能可以收到。那么我们基于这个去完善下我们的功能。
首先，我们的登录用户还是用了之前分享中分享的数据库相关呢，我们现在直接来调用，我们之前有从数据库获取用户，然后校验密码，密码成功我们才让登录成功。这个FastAPI 学习之路（四十六）WebSockets（三）登录后才可以聊天就已经分享了。
其次，我们去建立了链接，但是我们现在的链接都是没有针对发消息做管理呢，我们去扩展下我们之前的发消息的类。
class 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}) def disconnect(self, user: str, ws: WebSocket): # 关闭时 移除ws对象 self.active_connections.remove({"user": user, "ws": ws}) @staticmethod async def send_personal_message(message: str, ws: WebSocket): # 发送所有人消息 await ws.send_text(message) async def send_other_message(self, message: dict, user: str): # 发送个人消息 for connection in self.active_connections: if connection["user"] == user: await connection['ws'].send_json(message) async def broadcast(self, data: str): # 广播消息 for connection in self.active_connections: await connection['ws'].send_text(data)
增加了广播消息，和增加了发所有消息和针对个人发消息。我们现在实现的是针对所有人和广播。那么我们看下，我们的管理聊天链接的已经完善。
接下来我们就是去实现如何接受消息。
@app.websocket("/items/ws/{user}") async def websocket_endpoint( websocket: WebSocket, user:str, cookie_or_token: str = Depends(get_cookie_or_token), ): await manager.connect(user, websocket) await manager.broadcast(user + "进入聊天室") try: while True: data = await websocket.receive_text() await manager.send_personal_message(f"你说了: {data}", websocket) await manager.broadcast(f"用户:{user} 说: {data}") except WebSocketDisconnect as e: manager.disconnect(user, websocket) await manager.broadcast("用户{}离开".format(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>
其实我们就是在连接的时候增加了对应人的连接，那么我们看下，我们应该测试下效果。
首先我们看下，我们第一个用户已经进来了，我们看下第二个用户的登录。登录后，我们发现用户二进来了，我们看下用户1是否接受到用户二上线的消息了？
我们可以看到，当二完成了登录之后呢，用户一也接受到了用户二的消息
我们可以试下发送消息。
这样我们的聊天更加完善了。我们从简单的聊天室入手，慢慢去改造我们的代码，测试我们的代码。逐步去完善，慢慢去改造实现。
文章首发在公众号，欢迎关注。
