FastAPI WebSocket 分组发送Json数据
效果
用户1和 用户2 可以互相发送私信消息
用户1 2 3之间相当于一个群，可以发送广播消息
代码
FastAPI 服务端代码
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2020/8/19 13:38
# @Author : CoderCharm
# @File : main.py
# @Software: PyCharm
# @Github : github/CoderCharm
# @Email : wg_python@163.com
# @Desc :
"""
https://stackoverflow.com/questions/15219858/how-to-store-a-complex-object-in-redis-using-redis-py
obj = ExampleObject()
pickled_object = pickle.dumps(obj)
r.set('some_key', pickled_object)
unpacked_object = pickle.loads(r.get('some_key'))
obj == unpacked_object
typing.Dict[key_type, value_type]
"""
from typing import List, Dict
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
app = FastAPI()
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: dict, ws: WebSocket):
# 发送个人消息
await ws.send_json(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: dict):
# 广播消息
for connection in self.active_connections:
await connection['ws'].send_json(data)
manager = ConnectionManager()
@app.websocket("/ws/{user}")
async def websocket_endpoint(ws: WebSocket, user: str):
await manager.connect(user, ws)
await manager.broadcast({"user": user, "message": "进入聊天"})
try:
while True:
data = await ws.receive_json()
print(data, type(data))
send_user = data.get("send_user")
if send_user:
await manager.send_personal_message(data, ws)
await manager.send_other_message(data, send_user)
else:
await manager.broadcast({"user": user, "message": data['message']})
except WebSocketDisconnect:
manager.disconnect(user, ws)
await manager.broadcast({"user": user, "message": "离开"})
if __name__ == "__main__":
import uvicorn
# 官方推荐是用命令后启动 uvicorn main:app --host=127.0.0.1 --port=8010 --reload
uvicorn.run(app='main:app', host="127.0.0.1", port=8010, reload=True, debug=True)
客户端代码
由于只是demo， 所以代码都是固定的，有三份固定写的身份信息， 到时候客户端会有登录
用户一
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>聊天1</title>
</head>
<body>
<h1>User1 Chat</h1>
<form action="" onsubmit="sendMessage(event)">
<input type="text" id="messageText" autocomplete="off"/>
<button>Send</button>
</form>
<form action="" onsubmit="sendOtherMessage(event)">
<input type="text" id="messageOther" autocomplete="off"/>
<button>Send Other</button>
</form>
<ul id='messages'>
</ul>
<script>
let ws = new WebSocket("ws://127.0.0.1:8010/ws/user1");
ws.onmessage = function(event) {
let messages = document.getElementById('messages')
let message = document.createElement('li');
console.log(event.data, typeof (event.data), 2222)
let receiveJson = JSON.parse(event.data);
console.log(receiveJson, typeof (receiveJson), 333);
let content = document.createTextNode(`${receiveJson.user}-${receiveJson.message}`);
message.appendChild(content);
messages.appendChild(message)
};
function sendMessage(event) {
let input = document.getElementById("messageText");
let message = {message: input.value, user: "user1"};
let messageJson = JSON.stringify(message);
ws.send(messageJson);
input.value = '';
event.preventDefault()
}
function sendOtherMessage(event) {
let input = document.getElementById("messageOther");
let message = {message: input.value, user: "user1", send_user: "user2"};
let messageJson = JSON.stringify(message);
ws.send(messageJson);
input.value = '';
event.preventDefault()
}
</script>
</body>
</html>
用户二
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>聊天2</title>
</head>
<body>
<h1>User2 Chat</h1>
<form action="" onsubmit="sendMessage(event)">
<input type="text" id="messageText" autocomplete="off"/>
<button>Send</button>
</form>
<form action="" onsubmit="sendOtherMessage(event)">
<input type="text" id="messageOther" autocomplete="off"/>
<button>Send Other</button>
</form>
<ul id='messages'>
</ul>
<script>
let ws = new WebSocket("ws://127.0.0.1:8010/ws/user2");
ws.onmessage = function(event) {
let messages = document.getElementById('messages')
let message = document.createElement('li');
console.log(event.data, typeof (event.data), 2222)
let receiveJson = JSON.parse(event.data);
console.log(receiveJson, typeof (receiveJson), 333);
let content = document.createTextNode(`${receiveJson.user}-${receiveJson.message}`);
message.appendChild(content);
messages.appendChild(message)
};
function sendMessage(event) {
let input = document.getElementById("messageText")
let message = {message: input.value, user: "user2"}
let messageJson = JSON.stringify(message);
ws.send(messageJson);
input.value = '';
event.preventDefault()
}
function sendOtherMessage(event) {
let input = document.getElementById("messageOther");
let message = {message: input.value, user: "user2", send_user: "user1"};
let messageJson = JSON.stringify(message);
ws.send(messageJson);
input.value = '';
event.preventDefault()
}
</script>
</body>
</html>
用户三
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>聊天3</title>
</head>
<body>
<h1>User3 Chat</h1>
<form action="" onsubmit="sendMessage(event)">
<input type="text" id="messageText" autocomplete="off"/>
<button>Send</button>
</form>
<ul id='messages'>
</ul>
<script>
let ws = new WebSocket("ws://127.0.0.1:8010/ws/user3");
ws.onmessage = function(event) {
let messages = document.getElementById('messages')
let message = document.createElement('li');
console.log(event.data, typeof (event.data), 2222)
let receiveJson = JSON.parse(event.data);
console.log(receiveJson, typeof (receiveJson), 333);
let content = document.createTextNode(`${receiveJson.user}-${receiveJson.message}`);
message.appendChild(content);
messages.appendChild(message)
};
function sendMessage(event) {
let input = document.getElementById("messageText")
let message = {message: input.value, user: "user3"}
let messageJson = JSON.stringify(message);
ws.send(messageJson);
input.value = '';
event.preventDefault()
}
</script>
</body>
</html>
一些扩展知识
- 关于websocket的链接异常捕获
以下代码catch是捕获不到异常的。
try {
let ws = new WebSocket("ws://127.0.0.1:8010/ws/user1");
ws.onerror = function (error) {
console.log(error, 111);
};
}catch (e) {
console.log(e, 222)
}
关于解释：
- websocket 和 socket.io对比
https://stackoverflow.com/questions/10112178/differences-between-socket-io-and-websockets
总结
websocket 的基本使用原理就是这样了, 全双工的传输协议真的很方便。