# 17 | 强大的装饰器 你好,我是景霄。这节课,我们一起来学习装饰器。 装饰器一直以来都是Python中很有用、很经典的一个feature,在工程中的应用也十分广泛,比如日志、缓存等等的任务都会用到。然而,在平常工作生活中,我发现不少人,尤其是初学者,常常因为其相对复杂的表示,对装饰器望而生畏,认为它“too fancy to learn”,实际并不如此。 今天这节课,我会以前面所讲的函数、闭包为切入点,引出装饰器的概念、表达和基本用法,最后,再通过实际工程中的例子,让你再次加深理解。 接下来,让我们进入正文一起学习吧! ## 函数->装饰器 ### 函数核心回顾 引入装饰器之前,我们首先一起来复习一下,必须掌握的函数的几个核心概念。 第一点,我们要知道,在Python中,函数是一等公民(first-class citizen),函数也是对象。我们可以把函数赋予变量,比如下面这段代码: ``` def func(message): print('Got a message: {}'.format(message)) send_message = func send_message('hello world') # 输出 Got a message: hello world ``` 这个例子中,我们把函数func赋予了变量send\_message,这样之后你调用send\_message,就相当于是调用函数func()。 第二点,我们可以把函数当作参数,传入另一个函数中,比如下面这段代码: ``` 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 ``` 这个例子中,我们就把函数get\_message以参数的形式,传入了函数root\_call()中然后调用它。 第三点,我们可以在函数里定义函数,也就是函数的嵌套。这里我同样举了一个例子: ``` 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 ``` 这段代码中,我们在函数func()里又定义了新的函数get\_message(),调用后作为func()的返回值返回。 第四点,要知道,函数的返回值也可以是函数对象(闭包),比如下面这个例子: ``` 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 ``` 这里,函数func\_closure()的返回值是函数对象get\_message本身,之后,我们将其赋予变量send\_message,再调用send\_message(‘hello world’),最后输出了`'Got a message: hello world'`。 ### 简单的装饰器 简单的复习之后,我们接下来学习今天的新知识——装饰器。按照习惯,我们可以先来看一个装饰器的简单例子: ``` 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(),而内部函数wrapper()中又会调用原函数greet(),因此,最后调用greet()时,就会先打印`'wrapper of decorator'`,然后输出`'hello world'`。 这里的函数my\_decorator()就是一个装饰器,它把真正需要执行的函数greet()包裹在其中,并且改变了它的行为,但是原函数greet()不变。 事实上,上述代码在Python中有更简单、更优雅的表示: ``` 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)`语句,只不过更加简洁。因此,如果你的程序中有其它函数需要做类似的装饰,你只需在它们的上方加上`@decorator`就可以了,这样就大大提高了函数的重复利用和程序的可读性。 ### 带有参数的装饰器 你或许会想到,如果原函数greet()中,有参数需要传递给装饰器怎么办? 一个简单的办法,是可以在对应的装饰器函数wrapper()上,加上相应的参数,比如: ``` 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): ... ``` 事实上,通常情况下,我们会把`*args`和`**kwargs`,作为装饰器内部函数wrapper()的参数。`*args`和`**kwargs`,表示接受任意数量和类型的参数,因此装饰器就可以写成下面的形式: ``` def my_decorator(func): def wrapper(*args, **kwargs): print('wrapper of decorator') func(*args, **kwargs) return wrapper ``` ### 带有自定义参数的装饰器 其实,装饰器还有更大程度的灵活性。刚刚说了,装饰器可以接受原函数任意类型和数量的参数,除此之外,它还可以接受自己定义的参数。 举个例子,比如我想要定义一个参数,来表示装饰器内部函数被执行的次数,那么就可以写成下面这种形式: ``` 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.__name__ ## 输出 'wrapper' help(greet) # 输出 Help on function wrapper in module __main__: wrapper(*args, **kwargs) ``` 你会发现,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' ``` ### 类装饰器 前面我们主要讲了函数作为装饰器的用法,实际上,类也可以作为装饰器。类装饰器主要依赖于函数`__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,然后打印,并且调用原函数。因此,在我们第一次调用函数example()时,num\_calls的值是1,而在第二次调用时,它的值变成了2。 ### 装饰器的嵌套 回顾刚刚讲的例子,基本都是一个装饰器的情况,但实际上,Python也支持多个装饰器,比如写成下面这样的形式: ``` @decorator1 @decorator2 @decorator3 def func(): ... ``` 它的执行顺序从里到外,所以上面的语句也等效于下面这行代码: ``` 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) greet('hello world') # 输出 execute decorator1 execute decorator2 hello world ``` ## 装饰器用法实例 到此,装饰器的基本概念及用法我就讲完了,接下来,我将结合实际工作中的几个例子,带你加深对它的理解。 ### 身份认证 首先是最常见的身份认证的应用。这个很容易理解,举个最常见的例子,你登录微信,需要输入用户名密码,然后点击确认,这样,服务器端便会查询你的用户名是否存在、是否和密码匹配等等。如果认证通过,你就可以顺利登录;如果不通过,就抛出异常并提示你登录失败。 再比如一些网站,你不登录也可以浏览内容,但如果你想要发布文章或留言,在点击发布时,服务器端便会查询你是否登录。如果没有登录,就不允许这项操作等等。 我们来看一个大概的代码示例: ``` 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, ...) ... ``` 这段代码中,我们定义了装饰器authenticate;而函数post\_comment(),则表示发表用户对某篇文章的评论。每次调用这个函数前,都会先检查用户是否处于登录状态,如果是登录状态,则允许这项操作;如果没有登录,则不允许。 ### 日志记录 日志记录同样是很常见的一个案例。在实际工作中,如果你怀疑某些函数的耗时过长,导致整个系统的latency(延迟)增加,所以想在线上测试某些函数的执行时间,那么,装饰器就是一种很常用的手段。 我们通常用下面的方法来表示: ``` 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): ... ``` 这里,装饰器log\_execution\_time记录某个函数的运行时间,并返回其执行结果。如果你想计算任何函数的执行时间,在这个函数上方加上`@log_execution_time`即可。 ### 输入合理性检查 再来看今天要讲的第三个应用,输入合理性检查。 在大型公司的机器学习框架中,我们调用机器集群进行模型训练前,往往会用装饰器对其输入(往往是很长的JSON文件)进行合理性检查。这样就可以大大避免,输入不正确对机器造成的巨大开销。 它的写法往往是下面的格式: ``` import functools def validation_check(input): @functools.wraps(func) def wrapper(*args, **kwargs): ... # 检查输入是否合法 @validation_check def neural_network_training(param1, param2, ...): ... ``` 其实在工作中,很多情况下都会出现输入不合理的现象。因为我们调用的训练模型往往很复杂,输入的文件有成千上万行,很多时候确实也很难发现。 试想一下,如果没有输入的合理性检查,很容易出现“模型训练了好几个小时后,系统却报错说输入的一个参数不对,成果付之一炬”的现象。这样的“惨案”,大大减缓了开发效率,也对机器资源造成了巨大浪费。 ### 缓存 最后,我们来看缓存方面的应用。关于缓存装饰器的用法,其实十分常见,这里我以Python内置的LRU cache为例来说明(如果你不了解 [LRU cache](https://en.wikipedia.org/wiki/Cache_replacement_policies#Examples),可以点击链接自行查阅)。 LRU cache,在Python中的表示形式是`@lru_cache`。`@lru_cache`会缓存进程中的函数参数和结果,当缓存满了以后,会删除least recenly used 的数据。 正确使用缓存装饰器,往往能极大地提高程序运行效率。为什么呢?我举一个常见的例子来说明。 大型公司服务器端的代码中往往存在很多关于设备的检查,比如你使用的设备是安卓还是iPhone,版本号是多少。这其中的一个原因,就是一些新的feature,往往只在某些特定的手机系统或版本上才有(比如Android v200+)。 这样一来,我们通常使用缓存装饰器,来包裹这些检查函数,避免其被反复调用,进而提高程序运行效率,比如写成下面这样: ``` @lru_cache def check(param1, param2, ...) # 检查用户设备类型,版本号等等 ... ``` ## 总结 这节课,我们一起学习了装饰器的概念及用法。**所谓的装饰器,其实就是通过装饰器函数,来修改原函数的一些功能,使得原函数不需要修改。** > Decorators is to modify the behavior of the function through a wrapper so we don’t have to actually modify the function. 而实际工作中,装饰器通常运用在身份认证、日志记录、输入合理性检查以及缓存等多个领域中。合理使用装饰器,往往能极大地提高程序的可读性以及运行效率。 ## 思考题 那么,你平时工作中,通常会在哪些情况下使用装饰器呢?欢迎留言和我讨论,也欢迎你把这篇文章分享给你的同事、朋友,一起在交流中进步。