就是要摸鱼

python装饰器

发表于 分类于
本文字数： 5.1k 阅读时长 ≈ 7 分钟
本文介绍了python装饰器的原理，以及常见的打印日志的用法

python装饰器

Python 装饰器详解_python 带参数的装饰器-CSDN博客

对装饰器@wraps的解释（一看就懂）— 并对装饰器详解 - 交流_QQ_2240410488 - 博客园 (cnblogs.com)

1. 为什么需要装饰器呢？

Leader让小A写两个数字相加和相减的函数，小A很快就写完了：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def add(x, y):
    return x + y


def sub(x, y):
    return x - y


if __name__ == '__main__':
    result = add(1, 2)
    print(result)

    result = sub(5, 4)
    print(result)

# 输出：
# 3
# 1

Leader让小A添加上统计函数的运行时长的功能, 小A直接在调用函数时加上了时长的计算:

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
import time


def add(x, y):
    return x + y


def sub(x, y):
    return x - y


if __name__ == '__main__':
    start = time.time()
    result_1 = add(1, 2)
    end = time.time()
    print('result: %d' % result_1)
    print('time taken %f' % (end - start))

    start = time.time()
    result_2 = sub(5, 4)
    end = time.time()
    print('result: %d' % result_2)
    print('time taken %f' % (end - start))

# 输出:
# result: 3
# time taken 0.000000
# result: 1
# time taken 0.000000

Leader看了，说每次调用函数岂不是要写很多重复代码吗。小A进行了优化：

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
import time


def add(x, y):
    start = time.time()
    rv = x + y
    end = time.time()
    print('time taken %f' % (end - start))
    return rv


def sub(x, y):
    start = time.time()
    rv = x - y
    end = time.time()
    print('time taken %f' % (end - start))
    return rv


if __name__ == '__main__':
    result_1 = add(1, 2)
    print('result: %d' % result_1)

    result_2 = sub(5, 4)
    print('result: %d' % result_2)
    
# 输出
# time taken 0.000000
# result: 3
# time taken 0.000000
# result: 1

这种方法肯定比前一种要好。但是当我们有多个函数时，那么这似乎就不方便了。

小A又定义了一个计时的函数并包装其他函数，然后返回包装后的函数：

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
import time


def time_taken(func):
    def inner(*args, **kwargs):
        start = time.time()
        rv = func(*args, **kwargs)
        end = time.time()
        print('time taken %f' % (end - start))
        return rv

    return inner


def add(x, y):
    return x + y


def sub(x, y):
    return x - y


if __name__ == '__main__':
    add = time_taken(add)  # 将函数作为参数传给另一个函数
    result_1 = add(1, 2)
    print('result: %d' % result_1)

    sub = time_taken(sub)
    result_2 = sub(5, 4)
    print('result: %d' % result_2)

# 输出：
# time taken 0.000000
# result: 3
# time taken 0.000000
# result: 1

Leader说上面的解决方案以及非常接近装饰器的思想了，小A查了一下装饰器的用法，加入装饰器后代码果然变得很优雅。

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
import time


def time_taken(func):
    def inner(*args, **kwargs):
        start = time.time()
        rv = func(*args, **kwargs)
        end = time.time()
        print('time taken %f' % (end - start))
        return rv
    return inner


@time_taken
# @time_taken等价于add = time_taken(add)
def add(x, y):
    return x + y


@time_taken
def sub(x, y):
    return x - y


if __name__ == '__main__':
    result_1 = add(1, 2)
    print('result: %d' % result_1)

    result_2 = sub(5, 4)
    print('result: %d' % result_2)

# 输出：
# time taken 0.000000
# result: 3
# time taken 0.000000
# result: 1

2. 类装饰器

装饰器本质上是一个 Python 函数，它可以让其他函数在不需要做任何代码变动的前提下增加额外功能，装饰器的返回值也是一个函数对象。它经常用于有切面需求的场景，比如：插入日志、性能测试、事务处理、缓存、权限校验等场景。装饰器是解决这类问题的绝佳设计，有了装饰器，我们就可以抽离出大量与函数功能本身无关的雷同代码并继续重用。

2.1 不带参数的类装饰器

基于类装饰器的实现，必须实现__call__和__init__两个内置函数。
__init__ ：接收被装饰函数
__call__ ：实现装饰逻辑

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class logger(object):
	def __init__(self,func):
		self.func = func

	def __call__(self,*args,**kwargs):
		print 'the function {}() is running...'\
			.format(self.func.__name__)
		return self.func(*args,**kwargs)

@logger
def say(something):
	print 'say {}!'.format(something)

say('hello')

运行结果如下：

1
2
Copythe function say() is running...
say hello!
2.2 带参数的类装饰器

带参数和不带参数的类装饰器有很大的不同。
__init__ ：不再接收被装饰函数，而是接收传入参数
__call__ ：接收被装饰函数，实现装饰逻辑

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class logger(object):
	def __init__(self,level='INFO'):
		self.level = level

	def __call__(self,func):
		def wrapper(*args,**kwargs):
			print '{level}: the function {func} () is running...'\
				.format(level=self.level, func=func.__name__)

			func(*args,**kwargs)
		return wrapper

@logger(level='WARNING')
def say(something):
	print 'say {}!'.format(something)

say('hello')

运行结果如下：

1
2
CopyWARNING: the function say () is running...
say hello!
2.3 wraps语法糖

使用这个语法糖保留函数原来的属性，name，doc等属性

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
from functools import wraps


def logging(func):
    @wraps(func)
    def inner(*args, **kwargs):
        """logging.inner"""
        if type == 'debug':
            print('[DEBUG] logging')
        else:
            print('[INFO] logging')
        rv = func(*args, **kwargs)
        return rv
    return inner


@logging
def add(x, y):
    """Add x and y"""
    return x + y


if __name__ == '__main__':
    print(add.__name__)
    print(add.__doc__)

# 输出:
# add
# Add x and y

2.4 打印日志
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
# 在增加一层函数
from functools import wraps
import time
from random import randint

def record(output):
    def use_time(func):
        @wraps(func)
        def wrapper(*args,**kwargs):
            st_time = time.time()
            result = func(*args,**kwargs)
            end_time = time.time()
#             print(f'{func.__name__}函数use_time:{end_time-st_time}s')
            output(func.__name__, end_time-st_time)
        return wrapper
    return use_time

# 改装饰器的结果就可以自定义了，下面以print函数为例
@record(print)
def foo():
    time.sleep(randint(2,5))

for _ in range(3):
    foo()
>>>>
foo 3.000645875930786
foo 4.003818988800049
foo 2.0020666122436523

# 自己定义函数，输入日志文件
def write_log(name,content):
    with open('./time.log','a')as f:
        f.write(f'{name}耗时：{content}\r\n') # \r\n 换行
        
# 只需要将装饰器改为@record(write_log)
@record(write_log)
def foo():
    time.sleep(randint(2,5))

for _ in range(3):
    foo()