Zwlin's Blog

Python装饰器(一)

2020/01/03

给函数添加一个包装

给函数添加一个包装层(wrapper layer),以添加额外的处理,例如记录日志,计时统计等,可以通过自定义一个装饰器函数。举例如下:

 1import time
 2
 3from functools import wraps
 4
 5def timethis(func):
 6    '''
 7    Decorator that reports the execution time
 8    :param func:
 9    :return:
10    '''
11
12    @wraps(func)
13    def wrapper(*args, **kwargs):
14        start = time.time()
15        result = func(*args, **kwargs)
16        end = time.time()
17        print(func.__name__, end - start)
18        return result
19
20    return wrapper
21
22@timethis
23def countdown(cnt):
24    while cnt > 0:
25        cnt -= 1
26
27
28if __name__ == '__main__':
29    countdown(100000)
30    countdown(1000000)

装饰器就是这样一个函数,它可以接受一个函数作为输入并返回一个新的函数作为输出,当像这样编写代码时:

1@timethis
2def countdown(cnt):
3    ...

和单独执行下列步骤的效果是一样的:

1def countdown(cnt):
2    ...
3countdown = timethis(countdown)

同时,内建的装饰器比如@staticmethod,@classmethod,@property的工作方式也是一样的。

装饰器内部一般会涉及创建一个新的函数,利用*arg**kwargs来接受任意的参数。示例中的wrapper函数正是这样做的。在这个函数内部,我们需要调用原来的输入函数(即被包装的那个函数,它是装饰器的输入参数)并返回它的结果。但是,也可以添加任何想要添加的额外代码(例如计时处理)。这个新创建的wrapper函数会作为装饰器的结果返回,取代了原来的函数。

需要重点强调的是,装饰器一般来说不会修改调用签名(参数个数,类型,顺序等),也不会修改被包装函数返回的结果。 这里对*arg**kwargs的使用是为了确保可以接受任何形式的输入参数。装饰器的返回值几乎总是同调用func(*args,*kwargs)的结果一致。

编写装饰器时要保存函数的元数据

编写好了一个装饰器,但当他用在一个函数上时,一些重要的元数据比如函数名,文档字符串,函数注解以及调用签名都丢失了,因此,每当定义一个装饰器时,应该总是记得为底层的包装函数添加functools库中的@warps装饰器。示例如下:

 1import time
 2
 3from functools import wraps
 4
 5def timethis(func):
 6    '''
 7    Decorator that reports the execution time
 8    :param func:
 9    :return:
10    '''
11
12    @wraps(func)
13    def wrapper(*args, **kwargs):
14        start = time.time()
15        result = func(*args, **kwargs)
16        end = time.time()
17        print(func.__name__, end - start)
18        return result
19
20    return wrapper

对装饰器进行解包装

@wraps装饰器的一个重要特性就是它可以通过__warpped__属性来访问被包装的那个函数。例如,若果希望直接访问被包装的函数,这可以通过这样做:

1countdown.__wrapped__(10000)
2#or
3origin_countdown = countdown.__wrapped__
4origin_countdown()

这种技术只有在实现装饰器时利用了functools库中的@warps对元数据进行了适当的拷贝,或者直接设定了__warpped__属性时才有用。

同时,嵌套时的装饰器链,也是逐层解包的:

 1from functools import wraps
 2
 3
 4def decorator1(func):
 5    @wraps(func)
 6    def wrapper(*args, **kwargs):
 7        print("Decorator1")
 8        return func(*args, **kwargs)
 9
10    return wrapper
11
12
13def decorator2(func):
14    @wraps(func)
15    def wrapper(*args, **kwargs):
16        print("Decorator2")
17        return func(*args, **kwargs)
18
19    return wrapper
20
21
22@decorator1
23@decorator2
24def add(x, y):
25    return x + y
26
27
28if __name__ == '__main__':
29    print(add(2, 3))
30    print(add.__wrapped__(2, 3))
31    print(add.__wrapped__.__wrapped__(2, 3))
32
33# 输出
34Decorator1
35Decorator2
365
37Decorator2
385
395