Litestar框架中的WebSocket使用指南

侯霆垣

于 2025-06-06 09:11:31 发布

阅读量381

点赞数 4

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00267/article/details/148465859

Litestar框架中的WebSocket使用指南

litestar Production-ready, Light, Flexible and Extensible ASGI API framework | Effortlessly Build Performant APIs 项目地址: https://gitcode.com/gh_mirrors/li/litestar

WebSocket作为现代Web应用中的重要通信协议，在实时应用开发中扮演着关键角色。Litestar框架提供了多种灵活的方式来处理WebSocket连接，本文将深入解析这些方法及其适用场景。

WebSocket处理方式概览

Litestar框架提供了三种不同层级的WebSocket处理方式，满足从基础到高级的各种需求：

底层WebSocket路由处理器：提供对ASGI WebSocket接口的基本抽象
WebSocket监听器(WebsocketListener)：响应式、事件驱动的处理方式，支持完整序列化和DTO
WebSocket流(websocket_stream)：主动式、面向流的处理方式，同样支持完整序列化和DTO

选择哪种方式取决于应用场景：需要精细控制时选择底层接口，而高级接口则能简化开发流程。

WebSocket监听器详解

WebSocket监听器采用事件驱动模式，开发者只需关注业务逻辑，底层细节由框架处理。

基础示例

from litestar import Litestar
from litestar.handlers.websocket_handlers import websocket_listener

@websocket_listener("/")
async def handler(data: str) -> str:
    return data

app = Litestar([handler])

这个处理器会：

监听根路径("/")的连接
接收数据后自动传递给处理函数
返回数据会自动序列化并发送

数据接收处理

监听器支持多种数据类型接收，通过类型注解自动转换：

# 接收JSON数据
@websocket_listener("/json")
async def json_handler(data: dict) -> dict:
    return {"response": data}

# 接收文本数据
@websocket_listener("/text")
async def text_handler(data: str) -> str:
    return data.upper()

# 接收二进制数据
@websocket_listener("/bytes")
async def bytes_handler(data: bytes) -> bytes:
    return data[::-1]

数据发送机制

发送数据只需从处理函数返回即可，支持多种数据类型：

# 发送文本
@websocket_listener("/send-text")
async def send_text() -> str:
    return "Hello, WebSocket!"

# 发送二进制
@websocket_listener("/send-bytes")
async def send_bytes() -> bytes:
    return b"\x01\x02\x03"

# 发送复杂对象(自动序列化为JSON)
@websocket_listener("/send-complex")
async def send_complex() -> dict:
    return {"status": "ok", "data": [1, 2, 3]}

传输模式配置

WebSocket支持文本和二进制两种传输模式，可独立配置接收和发送模式：

# 配置接收模式
@websocket_listener("/binary-receive", receive_mode="binary")
async def binary_receive(data: bytes) -> bytes:
    return data

# 配置发送模式
@websocket_listener("/binary-send", send_mode="binary")
async def binary_send() -> bytes:
    return b"binary data"

高级特性

依赖注入

与常规路由处理器类似，WebSocket监听器也支持依赖注入：

from typing import Annotated
from litestar.params import Dependency

async def get_db_connection():
    # 模拟数据库连接
    print("建立连接")
    yield "db_connection"
    print("关闭连接")

@websocket_listener("/di")
async def di_handler(
    db: Annotated[str, Dependency(get_db_connection)]
) -> str:
    return f"使用连接: {db}"

直接访问WebSocket对象

需要底层控制时，可直接注入WebSocket实例：

@websocket_listener("/socket-access")
async def socket_access(
    data: str,
    socket: WebSocket
) -> None:
    await socket.send_text(f"ECHO: {data}")

类式监听器

对于复杂逻辑，可使用类式监听器：

from litestar.handlers import WebsocketListener

class ChatListener(WebsocketListener):
    path = "/chat"
    
    def on_connect(self) -> None:
        print("客户端连接")
    
    async def on_disconnect(self) -> None:
        print("客户端断开")
    
    async def on_receive(self, data: str) -> str:
        return f"收到: {data}"

WebSocket流处理

WebSocket流适合主动推送数据的场景：

import asyncio
from datetime import datetime

@websocket_stream("/time-feed")
async def time_stream() -> str:
    while True:
        yield str(datetime.now())
        await asyncio.sleep(0.5)