python asyncio 回调式代码封装为异步/等待式代码

应用场景

通过 python aiohttp 建立 websocket 连接时，需要通过 ws 发送 API 请求并拿到 API 返回值，并封装为result = await api(...)的形式。

开发环境

• python 3.9.4
• aiohttp 3.8.1

原理

使用 asyncio.Future 对象，在调用处等待 Future 对象的结果，在 ws 收到 API 的返回值时为 Future 设置结果(set_result)。

创建一个 Future 对象 future = loop.create_future()，使用 await asyncio.wait_for(future, 30, loop=loop)等待 Future 对象的结果。

在实际获取到返回值的地方为这个 Future 对象设置结果 future.set_result(result)。

实例

API 调用封装部分

import asyncio

# echo 和回调函数的映射
__CALLBACK = {}
# 回调识别码
ECHO = 0

class APITimeoutError(asyncio.TimeoutError):
    pass

def del_callback(echo: int):
    """删除callback"""
    if echo in __CALLBACK:
        del __CALLBACK[echo]

async def call_api(action: str, params: dict, ws) -> dict:
    # 调用API
    global ECHO
    loop = asyncio.get_running_loop()
    future = loop.create_future()
    ECHO += 1
    __CALLBACK[ECHO] = future
    # 这里发送的 echo 会在返回值里带上，
    loop.create_task(ws.send_json({
        'action': action,
        'params': params,
        'echo': ECHO
    }))
    try:
        # 等待 future 的结果
        return await asyncio.wait_for(future, 30, loop=loop)
    except asyncio.TimeoutError:
        # 超时则删除
        del_callback(ECHO)
        raise APITimeoutError

async def api_callback(result: dict, echo: int):
    if echo in __CALLBACK:
        # 从 __CALLBACK 中取出 echo 对应的回调函数
        __CALLBACK[echo].set_result(result)
        del_callback(echo)

ws 服务部分

@routes.get('/ws')
async def websocket_handler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    async for msg in ws:
        if msg.type == aiohttp.WSMsgType.TEXT:
            loop = asyncio.get_running_loop()
            # 此处不能使用await，否则会阻塞在等待 Future 对象结果处
            # await 不执行完毕不会继续往下执行
            # create_task 安排一个协程的运行，会继续往下执行
            # 详见 https://docs.python.org/zh-cn/3/library/asyncio-eventloop.html#asyncio.loop.create_task
            loop.create_task(on_message(msg.data, ws))

    return ws

async def on_message(message, ws):
    if message为API返回值:
        await api_callback(返回值, 返回值中带的echo)
    else:
        其它处理