1、简单的装饰器

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

这里的函数 my_decorator() 就是一个装饰器,它把真正需要执行的函数 greet() 包裹在其中,并且改变了它的行为,但是原函数 greet() 功能不变。

事实上,上述代码可以进一步简化:

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) 语句,只不过更加简洁。

2、带有参数的装饰器

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

如果另外还有一个函数,也需要使用 my_decorator() 装饰器,但是这个新函数有两个参数,又该怎么办?

@my_decorator
def celebrate(name, message):
    pass

通常情况下,我们会把 *args 和 **kwargs 作为装饰器内部函数 wrapper() 的参数。*args 和 **kwargs,表示接受任意数量和类型的参数,因此装饰器就可以写成下面的形式:

def my_decorator(func):
    def wrapper(*args, **kwargs):
        print('wrapper of decorator')
        func(*args, **kwargs)
    return wrapper

3、带有自定义参数的装饰器

装饰器可以接受原函数任意类型和数量的参数,除此以外,它还可以接受自定义的参数。例如,我们想要定义一个参数,表示装饰器内部函数被执行的次数,代码如下:

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

注意: 原函数还是原函数吗?

greet() 函数被装饰器装饰以后,它的元信息变了,元信息表明它不再是以前的那个 greet() 函数,而是被 wrapper() 函数取代了。

为了解决这个问题,我们通常使用内置的装饰器 @functools.wrap ,它会帮助保留原函数的元信息(也就是将原函数的元信息,拷贝到对应的装饰器函数里)。

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'

3、类装饰器

前面我们主要讲了函数作为装饰器的用法,实际上,类也可以作为装饰器。

类装饰器主要依赖于函数 \_\_call\_\_(),每当你调用一个类的实例时,函数 \_\_call\_\_() 就会被执行一次。

代码如下:

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

这里,我们定义了类 Count,初始化时传入原函数 func(),而 \_\_call\_\_() 函数表示让变量 num_calls 自增 1,然后打印,并且调用原函数。

装饰器的嵌套

上述讲的案例,基本上都是一个装饰器的情况,实际上,Python 也支持多个装饰器,比如写成如下形式:

@decorator1
@decorator2
@decorator3
def func():
    pass

它的执行顺序从里到外,所以上面的语句也等效于下面这行代码:

decorator1(decorator2(decorator3(func)))

这样,'hello world' 这个例子,就可以改写成下面的样子:

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)
    
# 输出
# execute decorator1
# execute decorator2
# hello world

4、装饰器用法实例

以上就是装饰器的基本概念以及用法,接下来,我将结合实际工作中的几个例子,带你加深对它的理解。

身份认证

首先是最常见的身份认证的应用。登录微信,需要输入用户名密码,服务器端便会查询你的用户名是否存在、密码是否匹配等等。如果认证通过,就可以顺利登录,如果认证失败,就抛出异常并提示登录失败。

再比如某些网站,即使你不登录也可以浏览内容,但是,如果想要发布文章或者留言,在点击发布时,服务器端便会查询你是否已经登录。如果没有登录,就不允许这项操作等等。

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, ...):
    pass

以上代码中,我们定义了装饰器 authenticate,而函数 post_comment(),则表示发表用户对谋篇文章的评论。每次调用这个函数之前,都会先检查用户是否处于登录状态,如果是登录状态,则允许这项操作;如果没有登录,则不允许这项操作。

日志记录

日志记录同样是很常见的一个案例。在实际工作中,如果你怀疑某些函数的耗时过长,导致整个系统的 latency 延迟增加,所以想在线上测试某些函数的执行时间,那么,装饰器就是一种很常见的手段。

代码如下:

import time
import functools

def log_execution_time(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start = time.pref_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):
    pass

装饰器 log_execution_time 记录某个函数的运行时间,并返回其执行结果。如果你想计算任何函数的执行时间,在这个函数上面加上 @log_execution_time 即可。

输入合理性检查

在大型公司的机器学习框架中,我们调用机器集群进行模型训练前,往往会用装饰器对其输入(往往是很长的 json 文件)进行合理性检查。这样就可以大大避免,输入不正确时对机器造成的巨大开销。

代码如下:

import functools

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

其实在工作中,很多情况下都会出现不合理的现象。因为我们调用的训练模型往往很复杂,输入的文件有成千上万行,很多时候确实也很难发现。

如果没有输入的合理性检查,很容易出现“模型训练了好几个小时后,系统却报错说输入的一个参数不对,成果付之一炬”的现象。这种常见的错误,会严重影响开发效率,也对机器资源造成了巨大浪费。

缓存

关于缓存装饰器的用法,其实十分常见,这里以 Python 内置的 LRU cache 为例来说明。

LRU cache,在 Python 中的表示形式是: @lru_cache,@lru_cache 会缓存进程中的函数参数和结果,当缓存满了以后,会删除 least recenly used 的数据。

正确使用缓存装饰器,往往能极大地提高程序运行效率。

大型公司服务器端的代码中往往存在很多关于设备的检查,比如你使用的设备是安卓还是 iPhone ,版本号是多少。这其中的一个原因就是一些新的 feature,往往只在某些特定的手机系统或版本上才有(比如 Android v200+)。

这样一来,我们通常使用缓存装饰器,来包裹这些检查函数,避免其被反复调用,进而提高程序运行效率,比如写成下面形式:

@lru_cache
def check(param1, param2, ...):   # 检查用户设备类型,版本号等等
    pass

思考

Python 的装饰器的应用场景有点像 AOP 的应用场景,把一些常用的业务逻辑分离,提高程序可重用性,降低耦合度,提高开发效率。

Last modification:November 19, 2021
如果觉得我的文章对你有用,请随意赞赏