fastapi 实现websocket协议

  fastapi是一个高性能异步web框架，并且支持websocket协议。这是一个比较新的框架，但github上的Star数挺多，快追上flask了。

 

1、websocket服务端的简单实现：

定义一个websocket API和 http 类似，使用@app.websockert装饰器，参数为路由。

@app.websocket('/test')
async def websocket_test(websocket: WebSocket):
    # 服务器接受客户端的WebSocket连接请求。
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"服务器返回：{data}")
    # 客户端断开连接，捕获WebSocketDisconnect异常
    except WebSocketDisconnect:
        pass

fastapi官方推荐使用uvicorn 作为ASGI服务器使用：

from fastapi import FastAPI
from fastapi import WebSocket,WebSocketDisconnect
import uvicorn

app = FastAPI()

@app.websocket('/test')
async def websocket_test(websocket: WebSocket):
    # 服务器接受客户端的WebSocket连接请求。
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"服务器返回：{data}")
    # 客户端断开连接，捕获WebSocketDisconnect异常
    except WebSocketDisconnect:
        pass
    
if __name__ == "__main__":
    uvicorn.run(app='main:app',host="0.0.0.0",port=8000)

 

客户端代码：

import websockets, asyncio

async def ws_client():
    async with websockets.connect('ws://127.0.0.1:8000/ai/test') as ws:
        await ws.send('this is a test')
        response = await ws.recv()
        print(response)

if __name__ == '__main__':
    asyncio.run(ws_client())

 

2、多人聊天

以下是官网的websocket服务端demo，实现多人聊天。自定义一个ConnectionManager类，和已建立连接的ws进行通信，我加了一些注释，方便查看。

from fastapi import FastAPI
import uvicorn
from fastapi import WebSocket
from typing import List

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)
           
manager = ConnectionManager()

@app.websocket('/object/{client_id}')
async def websocket_endpoint(websocket: WebSocket, client_id: int):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            # 返回给当前ws发送消息的客户端
            await manager.send_personal_message(f"You wrote: {data}", websocket)
            # 返回给所有已连接状态的客户端
            await manager.broadcast(f"Client #{client_id} says: {data}")
    except WebSocketDisconnect:
        # 断开ws连接
        manager.disconnect(websocket)
        # 告诉剩余还在线客户端，Client是{client_id}的ws断开了
        await manager.broadcast(f"Client #{client_id} left the chat")
        
        
if __name__ == "__main__":
    uvicorn.run(app='main:app',host="0.0.0.0",log_level="debug",port=8000)