如何在FastAPI中自定义WebSocket对象？类比自定义Request的实现方式

要实现自定义WebSocket对象并替换Starlette默认实例，你可以通过子类化WebSocketRoute来完成，类似HTTP路由中自定义APIRoute的思路，具体步骤如下：

1. 定义自定义WebSocket类 #

继承Starlette原生的WebSocket，扩展构造逻辑和自定义方法：

from starlette.websockets import WebSocket as StarletteWebSocket

class CustomWebSocket(StarletteWebSocket):
    def __init__(self, scope, receive, send):
        # 添加自定义初始化属性
        self.user_context = {}
        super().__init__(scope, receive, send)
    
    # 扩展自定义方法
    async def send_formatted_message(self, content: str):
        """发送带前缀的格式化文本消息"""
        await self.send_text(f"[系统通知] {content}")

2. 子类化WebSocketRoute替换实例 #

重写WebSocketRoute的get_application方法，在ASGI应用链中注入自定义WebSocket实例：

from starlette.routing import WebSocketRoute

class CustomWebSocketRoute(WebSocketRoute):
    def get_application(self):
        # 获取原路由的ASGI应用
        original_app = super().get_application()
        
        async def wrapper(scope, receive, send):
            # 创建自定义WebSocket实例
            custom_ws = CustomWebSocket(scope, receive, send)
            # 将自定义实例的参数传入原应用
            await original_app(custom_ws.scope, custom_ws.receive, custom_ws.send)
        
        return wrapper

3. 在FastAPI中使用自定义路由 #

在WebSocket路由装饰器中指定route_class参数，即可让端点接收自定义WebSocket对象：

from fastapi import FastAPI

app = FastAPI()

@app.websocket("/ws", route_class=CustomWebSocketRoute)
async def websocket_handler(websocket: CustomWebSocket):
    await websocket.accept()
    # 使用自定义方法发送消息
    await websocket.send_formatted_message("连接成功！")
    
    while True:
        data = await websocket.receive_text()
        # 访问自定义属性
        websocket.user_context["last_msg"] = data
        await websocket.send_text(f"收到消息：{data}")

原理说明 #

Starlette的WebSocketRoute通过get_application返回处理WebSocket连接的ASGI应用函数，我们通过包装这个函数，替换掉原生的WebSocket实例，让端点函数最终拿到的是自定义扩展后的对象。

内容的提问来源于stack exchange，提问作者acw