Python Websockets如何向client_handler传递Queue实现跨协程数据处理

你可以通过以下两种方式实现参数传递，无需使用全局变量，也不需要将client_handler定义为局部嵌套函数：

方法1：使用functools.partial绑定预定义参数 #

functools.partial可以为指定函数提前绑定部分参数，生成一个参数数量匹配websockets要求的新handler函数，操作如下：

1. 给原client_handler新增队列参数
2. 在启动服务时用partial将队列绑定到handler中

修改后的完整代码如下：

import asyncio
import functools
import websockets

# 新增q作为第一个参数
async def client_handler(q, websocket, path):
    print(f"Connected with path '{path}'")

    async for msg_rx in websocket:
        if not msg_rx:
            break
        print(f"RX: {msg_rx }")
        # 直接写入队列即可
        await q.put(msg_rx)

    print(f"Disconnected from Path '{path}'")


async def task_ws_server(q):
    # 用partial绑定队列参数，生成符合websockets要求的handler
    handler = functools.partial(client_handler, q)
    async with websockets.serve(handler, '127.0.0.1', 5001):
        await asyncio.Future()  # run forever

async def task_consumer(q):
    while True:
        data = await q.get()
        # 可替换为存储到文件、转发到其他服务等业务逻辑
        print(data)
        q.task_done()


async def main():
    q = asyncio.Queue()
    consumer = asyncio.create_task(task_consumer(q))
    await asyncio.gather(*[
        asyncio.create_task(task_ws_server(q)),
    ])
    await q.join()
    consumer.cancel()


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

方法2：使用工厂函数生成handler #

如果需要封装更多上下文参数，也可以用全局的工厂函数生成带队列上下文的handler，核心实现逻辑如下：

# 全局工厂函数，接收队列参数返回对应的handler
def get_client_handler(q):
    async def client_handler(websocket, path):
        print(f"Connected with path '{path}'")
        async for msg_rx in websocket:
            if not msg_rx:
                break
            print(f"RX: {msg_rx }")
            await q.put(msg_rx)
        print(f"Disconnected from Path '{path}'")
    return client_handler

# task_ws_server中调用工厂函数生成handler即可
async def task_ws_server(q):
    async with websockets.serve(get_client_handler(q), '127.0.0.1', 5001):
        await asyncio.Future()  # run forever

两种方案都不会将队列暴露到全局作用域，也不需要在其他函数内部定义client_handler，完全匹配你的需求。

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