一日一技：装饰器如何装饰异步函数

青南

发布于 2020-02-19 15:56:56

3.3K0

发布于 2020-02-19 15:56:56

在 Python 开发的工程中，我们常常使用装饰器来优化代码，例如一个打印日志的装饰器：

import requests
import datetime


def time_log(func):
    def wrap(*args, **kwargs):
        print(f'现在时间：{datetime.datetime.now()}，程序开始运行')
        result = func(*args, **kwargs)
        print(f'现在时间：{datetime.datetime.now()}，运行结束')
        return result
    return wrap


@time_log
def get(n):
    resp = requests.get(f'http://httpbin.org/delay/{n}').text
    print(resp)


get(3)

运行效果如下图所示：

但如果get是一个异步函数，这个装饰器就会出问题：

import pprint
import aiohttp
import datetime
import asyncio


def time_log(func):
    def wrap(*args, **kwargs):
        print(f'现在时间：{datetime.datetime.now()}，程序开始运行')
        result = func(*args, **kwargs)
        print(f'现在时间：{datetime.datetime.now()}，运行结束')
        return result
    return wrap

@time_log
async def get(n):
    async with aiohttp.ClientSession() as client:
        resp = await client.get(f'http://httpbin.org/delay/{n}')
        result = await resp.json()
    pprint.pprint(result, indent=2)


asyncio.run(get(3))

运行效果如下图所示：

可以看到，print(f'现在时间：{datetime.datetime.now()}，程序开始运行')与print(f'现在时间：{datetime.datetime.now()}，运行结束')几乎同时被打印出来，然后才是请求网络。这并不是我们需要实现的效果。我们想要的是先打印前一行，然后请求网络，再打印后一行。

为了解决这个问题，我们需要把time_log装饰器中的wrap也定义成异步函数：

def time_log(func):
    async def wrap(*args, **kwargs):
        print(f'现在时间：{datetime.datetime.now()}，程序开始运行')
        result = await func(*args, **kwargs)
        print(f'现在时间：{datetime.datetime.now()}，运行结束')
        return result
    return wrap

同时await func(*args, **kwargs)。这样就能保证代码的执行顺序：