如何为Python websockets的serve函数传递额外参数？

你遇到的错误是因为websockets.serve会将除处理函数、主机、端口之外的关键字参数传递给底层的asyncio.create_server，而后者不支持send_queue这类自定义参数。要给处理函数传递额外参数，有两种常用方案：

方案一：使用functools.partial绑定参数 #

通过functools.partial将队列参数预先绑定到处理函数上，生成一个符合serve要求签名的新函数：

import asyncio
import websockets
from functools import partial

async def receive_websocket_message(websocket, send_queue: asyncio.Queue, receive_queue: asyncio.Queue):
    async for message in websocket:
        # 向内部客户端发送消息
        send_queue.put_nowait(message)
        # 等待响应
        response = await receive_queue.get()
        # 向WebSocket发送方返回响应
        await websocket.send(response)

async def websocket_server(send_queue: asyncio.Queue, receive_queue: asyncio.Queue):
    # 绑定额外参数到处理函数
    handler = partial(receive_websocket_message, send_queue=send_queue, receive_queue=receive_queue)
    async with websockets.serve(handler, "localhost", 8765):
        await asyncio.Future()  # 永久运行

方案二：使用闭包捕获参数 #

在websocket_server内部定义处理函数，直接捕获外部作用域的队列变量：

import asyncio
import websockets

async def websocket_server(send_queue: asyncio.Queue, receive_queue: asyncio.Queue):
    async def receive_websocket_message(websocket):
        async for message in websocket:
            send_queue.put_nowait(message)
            response = await receive_queue.get()
            await websocket.send(response)
    
    async with websockets.serve(receive_websocket_message, "localhost", 8765):
        await asyncio.Future()  # 永久运行

两种方案都能让处理函数拿到需要的队列参数，同时满足websockets.serve对处理函数签名的要求。

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