给异步函数加上装饰器

前言

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

同步函数使用装饰器

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结束

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

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

总结

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

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