Python装饰器

发表于 | 更新于 | 分类于 ， ， | 阅读次数：

Python装饰器

本文为景霄-Python核心技术与实战的学习笔记，如要查看完整内容请点击链接。

从函数到装饰器

回顾函数

在Python中，函数也是对象，因而我们可以将函数赋予变量：

1
2
3
4
5
6
7
8
def func(message):
    print('Got a message: {}'.format(message))
    
send_message = func
send_message('hello world')

# 输出
Got a message: hello world

也可以将函数作为参数，传入另一个函数：

1
2
3
4
5
6
7
8
9
10
11
def get_message(message):
    return 'Got a message: ' + message


def root_call(func, message):
    print(func(message))
    
root_call(get_message, 'hello world')

# 输出
Got a message: hello world

也可以进行函数嵌套：

1
2
3
4
5
6
7
8
9
def func(message):
    def get_message(message):
        print('Got a message: {}'.format(message))
    return get_message(message)

func('hello world')

# 输出
Got a message: hello world

如果在上述嵌套中，返回的对象是函数，那么就构成了闭包：

1
2
3
4
5
6
7
8
9
10
def func_closure():
    def get_message(message):
        print('Got a message: {}'.format(message))
    return get_message

send_message = func_closure()
send_message('hello world')

# 输出
Got a message: hello world

装饰器

一个较为简单的自定义装饰器实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def my_decorator(func):
    def wrapper():
        print('wrapper of decorator')
        func()
    return wrapper

def greet():
    print('hello world')

greet = my_decorator(greet)
greet()

# 输出
wrapper of decorator
hello world

变量greet指向内部函数wrapper()，内部函数包裹原函数greet()。这里的my_decorator()就是一个装饰器，它将真正需要执行的函数greet()包裹在内，同时改变其行为，但不改变原函数。

所谓装饰器就是通过装饰器函数来修改（添加）原函数的一些功能，而不需要直接对原始函数进行修改。

Decorators is to modify the behavior of the function through a wrapper so we don’t have to actually modify the function.

在上述代码中，我们自己手动将greet()函数传入my_decorator()函数中，在Python中有着更为简洁的实现：

1
2
3
4
5
6
7
8
9
10
11
def my_decorator(func):
    def wrapper():
        print('wrapper of decorator')
        func()
    return wrapper

@my_decorator
def greet():
    print('hello world')

greet()

@被称为语法糖，@my_decorator相当于greet = my_decorator(greet)语句。

带有参数的装饰器

如果原函数存在参数，那么如何定义装饰器？

直接在对应的装饰器函数wrapper()上添加相应的参数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def my_decorator(func):
    def wrapper(message):
        print('wrapper of decorator')
        func(message)
    return wrapper


@my_decorator
def greet(message):
    print(message)


greet('hello world')
# 输出
wrapper of decorator
hello world

但是如果想把这个装饰器用于另一个具有多个参数的函数时，上述做法就不适用。
此时就可以使用*args和**kwargs，如下所示：

1
2
3
4
5
def my_decorator(func):
    def wrapper(*args, **kwargs):
        print('wrapper of decorator')
        func(*args, **kwargs)
    return wrapper

带有自定义参数的装饰器

装饰器除了可以接受原函数任意类型和数量的参数外，还可以定义自己接收的参数。

例如，使用一个参数表示装饰器内部函数被执行的次数：

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
def repeat(num):
    def my_decorator(func):
        def wrapper(*args, **kwargs):
            for i in range(num):
                print('wrapper of decorator')
                func(*args, **kwargs)
        return wrapper
    return my_decorator


@repeat(4)
def greet(message):
    print(message)

greet('hello world')

# 输出：
wrapper of decorator
hello world
wrapper of decorator
hello world
wrapper of decorator
hello world
wrapper of decorator
hello world

相当于在之前的my_decorator()之外再添加了一层函数，并且在定义装饰器时首先生成了外层函数的对象，返回my_decorator()。

原函数发生变化了吗？

对于上述例子，打印原函数的元信息：

1
2
3
4
5
6
7
8
9
greet.__name__
## 输出
'wrapper'

help(greet)
# 输出
Help on function wrapper in module __main__:

wrapper(*args, **kwargs)

发现，经过装饰器装饰后，原函数的信息发生了变化，变成了wrapper()函数。

此时，可以使用内置装饰器@functools.wrap来保留原函数的元信息：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

import functools

def my_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print('wrapper of decorator')
        func(*args, **kwargs)
    return wrapper
    
@my_decorator
def greet(message):
    print(message)

greet.__name__

# 输出
'greet'

类装饰器

除了函数可以被用作装饰器外，类也可以被用作装饰器。类装饰器主要通过函数__call__()实现：

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

class Count:
    def __init__(self, func):
        self.func = func
        self.num_calls = 0

    def __call__(self, *args, **kwargs):
        self.num_calls += 1
        print('num of calls is: {}'.format(self.num_calls))
        return self.func(*args, **kwargs)

@Count
def example():
    print("hello world")

example()

# 输出
num of calls is: 1
hello world

example()

# 输出
num of calls is: 2
hello world

...

装饰器的嵌套

Python支持多个装饰器的嵌套：

1
2
3
4
5
@decorator1
@decorator2
@decorator3
def func():
    ...

上述代码从里向外执行：

1
decorator1(decorator2(decorator3(func)))

多层嵌套例子

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

import functools

def my_decorator1(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print('execute decorator1')
        func(*args, **kwargs)
    return wrapper


def my_decorator2(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print('execute decorator2')
        func(*args, **kwargs)
    return wrapper


@my_decorator1
@my_decorator2
def greet(message):
    print(message)


greet('hello world')

# 输出
execute decorator1
execute decorator2
hello world

装饰器的用法

我们可以将在执行一个函数之前需要做的事情交给装饰器去做，同时不改变原始函数的行为。

身份认证

在进行某个操作前，常需要首先进行身份认证，此时就可以使用装饰器完成，使用装饰器来完成身份验证的任务：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

import functools

def authenticate(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        request = args[0]
        if check_user_logged_in(request): # 如果用户处于登录状态
            return func(*args, **kwargs) # 执行函数post_comment() 
        else:
            raise Exception('Authentication failed')
    return wrapper
    
@authenticate
def post_comment(request, ...)
    ...

这样，每次在调用post_comment()之前都会调用认证步骤。

日志记录

想要统计某段代码的运行耗时，就可以使用装饰器实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

import time
import functools

def log_execution_time(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        res = func(*args, **kwargs)
        end = time.perf_counter()
        print('{} took {} ms'.format(func.__name__, (end - start) * 1000))
        return res
    return wrapper
    
@log_execution_time
def calculate_similarity(items):
    ...

输入合理性检查

在大型公司的机器学习框架中，在调用机器进行模型训练前，常使用装饰器对器输入进行合理性检查，以避免输入错误带来的额外开销：

1
2
3
4
5
6
7
8
9
10
11

import functools

def validation_check(input):
    @functools.wraps(func)
    def wrapper(*args, **kwargs): 
        ... # 检查输入是否合法
    
@validation_check
def neural_network_training(param1, param2, ...):
    ...

缓存

LRU cache在Python中表示为@lru_cache，该装饰器会缓存进程中的函数参数和结果，当缓存满后，删除最近最久未使用的数据。

在实际工程中，可以将一些反复调用的函数使用@lru_cache装饰器进行装饰，这样，当调用相同相同的函数对象时便会直接从缓存中获得结果，而不用再次调用函数。

1
2
3
@lru_cache
def check(param1, param2, ...) # 检查用户设备类型，版本号等等
    ...