给异步函数加上装饰器

发表于 更新于 分类于

前言

平时使用装饰器的过程中，大多数情况都是用装饰器来装饰同步函数，如果有需要装饰异步函数的需求就不能像以前一样使用了，这篇文章来学习下如何使用装饰器来装饰异步函数。

同步函数使用装饰器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import requests
import datetime


def log(name='default'):
    def run(func):
        def a(*args, **kwargs):
            print(f'{str(datetime.datetime.now())[:-4]}，{name}运行')
            result = func(*args, **kwargs)
            print(f'{str(datetime.datetime.now())[:-4]}，{name}结束')
            return result

        return a

    return run


@log('get_message')
def demo():
    resp = requests.get('http://httpbin.org/ip').text
    print(resp)


if __name__ == '__main__':
    demo()

使用一个带参数的装饰器来装饰一个获取 IP 地址的函数，在函数运行前后打印时间。

1
2
3
4
5
6
2020-11-10 14:05:36.65，get_message运行
{
  "origin": "36.106.224.35"
}

2020-11-10 14:05:37.10，get_message结束

运行函数，可以看到装饰器正常运行，输出了函数结果和开始结束时间。

如果需要装饰的是一个异步函数呢？

异步函数使用装饰器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import asyncio

import datetime
import aiohttp


def log(name='default'):
    def run(func):
        def a(*args, **kwargs):
            print(f'{str(datetime.datetime.now())[:-4]}，{name}运行')
            result = func(*args, **kwargs)
            print(f'{str(datetime.datetime.now())[:-4]}，{name}结束')
            return result

        return a

    return run


@log('get_message')
async def demo():
    async with aiohttp.ClientSession() as session:
        resp = await session.get('http://httpbin.org/ip')
        print(await resp.text())


if __name__ == '__main__':
    asyncio.run(demo())
1
2
3
4
5
2020-11-10 14:16:48.95，get_message运行
2020-11-10 14:16:48.95，get_message结束
{
  "origin": "36.106.224.35"
}

运行以上代码，发现一瞬间就会打印出两条日志，时间时一样的，之后才打印请求的结果，这肯定是不正常的，异步函数装饰器不能直接使用。

那么需要怎么修改呢？其实只需要将装饰器内部的函数修改为异步函数即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import asyncio

import datetime
import aiohttp


def log(name='default'):
    def run(func):
        async def a(*args, **kwargs):
            print(f'{str(datetime.datetime.now())[:-4]}，{name}运行')
            result = await func(*args, **kwargs)
            print(f'{str(datetime.datetime.now())[:-4]}，{name}结束')
            return result

        return a

    return run


@log('get_message')
async def demo():
    async with aiohttp.ClientSession() as session:
        resp = await session.get('http://httpbin.org/ip')
        print(await resp.text())


if __name__ == '__main__':
    asyncio.run(demo())
1
2
3
4
5
6
2020-11-10 14:31:09.89，get_message运行
{
  "origin": "36.106.224.35"
}

2020-11-10 14:31:10.89，get_message结束

只需要将最内层的函数修改为异步即可，外部的两个函数依然是同步的，因为在 demo 函数还没有执行时，外部的两个函数就已经执行完了，不需要改成异步。

通用装饰器

基于上面的原理，可以将装饰器改为通用装饰器，在函数中判断传入的是普通函数还是异步函数，对应的返回不同的函数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import asyncio

import datetime
import aiohttp
from asyncio.coroutines import iscoroutinefunction

import requests


def log(name='default'):
    def run(func):
        async def a_async(*args, **kwargs):
            print(f'{str(datetime.datetime.now())[:-4]}，{name}运行')
            result = await func(*args, **kwargs)
            print(f'{str(datetime.datetime.now())[:-4]}，{name}结束')
            return result

        def a(*args, **kwargs):
            print(f'{str(datetime.datetime.now())[:-4]}，{name}运行')
            result = func(*args, **kwargs)
            print(f'{str(datetime.datetime.now())[:-4]}，{name}结束')
            return result

        if iscoroutinefunction(func):
            return a_async
        else:
            return a

    return run


@log('get_message')
async def demo_async():
    async with aiohttp.ClientSession() as session:
        resp = await session.get('http://httpbin.org/ip')
        print(await resp.text())


@log('get_message')
def demo():
    print(requests.get('http://httpbin.org/ip').json())


if __name__ == '__main__':
    asyncio.run(demo_async())
    print('-' * 20)
    demo()
1
2
3
4
5
6
7
8
9
10
2020-11-10 14:41:53.32，get_message运行
{
  "origin": "36.106.224.35"
}

2020-11-10 14:41:53.79，get_message结束
--------------------
2020-11-10 14:41:53.79，get_message运行
{'origin': '36.106.224.35'}
2020-11-10 14:41:54.26，get_message结束

使用 asyncio 的 iscoroutinefunction 函数来判断传入的函数类型，然后返回不同的装饰器函数。

可以看到，现在的装饰器已经变的通用了，异步和同步函数都可以使用同一个装饰器来装饰。

总结

python 的装饰器在打印日志，捕获异常等场合还是很有用的，这里写了一个简单的通用装饰器，可以装饰同步和异步的函数，不需要写重复的代码了。

本文章首发于个人博客 LLLibra146’s blog
本文作者：LLLibra146
更多文章请关注：qrcode
版权声明：本博客所有文章除特别声明外，均采用 © BY-NC-ND 许可协议。非商用转载请注明出处！严禁商业转载！
本文链接：
https://blog.d77.xyz/archives/e47a7f3c.html