0%

浅谈Python装饰器

装饰器

前言

Python 有一个有趣的功能,称为装饰器,它可以向现有代码添加功能

谈装饰器前,首先要明白一件事,Python 中的函数可以像普通变量一样当做参数传递给另外一个函数

事实上,Python 中的一切都是对象,函数也不例外,只不过其附带属性

先看一个例子

1
2
3
4
5
6
7
8
def foo(func):
func()

def bar():
print("I'm a function.")

if __name__ == "__main__":
foo(bar)

示例中,像foo这样能把函数作为参数的函数叫做高阶函数(higher order functions)

函数和方法都是可调用的,实际上,任何实现特殊__call __()方法的对象都是可调用的。因此,从基本的意义上讲,装饰器是可调用的,返回可调用的对象

第一个装饰器

装饰器接受一个函数,添加一些功能并返回它

我们实现一个函数,它接收一个函数,并且计算接收的函数执行了多久

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

def timer(func):
def wrapper():
start = time.time()
func()
end = time.time()
print('函数运行了', end-start, '秒')
return wrapper

def fun():
time.sleep(1)

if __name__ == '__main__':
instance = timer(fun)
instance()
# 函数运行了 1.0004889965057373 秒

这样是实现了计算函数运行时间的功能,但是很明显该操作比较麻烦,Python 提供的装饰器语法糖@就带来了极大的便利

我们这样修改代码,顺便打印函数的名称__name__

1
2
3
4
5
6
7
8
9
10
11
# ...

@timer
def fun():
time.sleep(1)

if __name__ == '__main__':
fun()
print(fun.__name__)
# 函数运行了 1.012272596359253 秒
wrapper

很明显,效果是等价的,但是函数fun()的名称却不是fun而是wrapper,其实没有错,因为装饰器返回的是wrapper函数,自然函数的名称也是wrapper

函数functools.wraps

对于上述问题,一个简单的解决办法就是我们可以添加一个代码块,把func的名称赋值给wrapper

1
2
3
4
5
def timer(func):
def wrapper():
wrapper.__name__ = func.__name__
# ...
return wrapper

这样再次执行输出的结果就是

1
2
函数运行了 1.0022237300872803 秒
fun

不过要是函数fun()中还有其他信息,如私有变量、文档说明等,这样的操作就太繁杂了

很棒的是,Python 内置的functools.wraps函数就帮我们解决了这样的事情,我们可以这样修改代码

1
2
3
4
5
6
7
8
9
10
from functools import wraps

def timer(func):
@wraps(func)
def wrapper():
start = time.time()
func()
end = time.time()
print('函数运行了', end-start, '秒')
return wrapper

再次执行,输出结果如下

1
2
函数运行了 1.0015637874603271 秒
fun

装饰带参函数

到此,我们学会了装饰器最基本的用法,不过现在装饰器装饰的是一个无参函数,若被装饰的函数有参数,其实也不难解决

我们可以写一个类似调试功能的装饰器,在调用函数时,返回函数名、参数及其返回值

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

def debug(func):
"""Print the function signature and return value"""
@functools.wraps(func)
def wrapper_debug(*args, **kwargs):
args_repr = [repr(a) for a in args]
kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
signature = ", ".join(args_repr + kwargs_repr)
print(f"Calling {func.__name__}({signature})")
value = func(*args, **kwargs)
print(f"{func.__name__!r} returned {value!r}")
return value
return wrapper_debug

@debug
def make_greeting(name, age=None):
if age is None:
return f"Howdy {name}!"
else:
return f"Whoa {name}! {age} already, you are growing up!"

if __name__ == '__main__':
print(make_greeting("Nopech", age=20))
# Calling make_greeting('Nopech', age=20)
# 'make_greeting' returned 'Whoa Nopech! 20 already, you are growing up!'
# Whoa Nopech! 20 already, you are growing up!

其中,因为我们不知道有多少个参数,所以可以用*args;当然还可能会有关键字参数,一律用**kwargs来代表

带参数的装饰器

其实,不光传入的函数可以有参数,装饰器本身也可以带参数

假设有这样一个需求,实现一个装饰器,能够指定日志的等级,并打印消息

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def logger(level):
def decorator(func):
def wrapper(*args, **kwargs):
print("[{level}]The function({name}) is running.".format(level=level, name=func.__name__))
return func(*args, **kwargs)
return wrapper
return decorator

@logger(level="INFO")
def fun(name='foo'):
print("I am %s." % name)

if __name__ == '__main__':
fun('a Function')
# [INFO]The function(fun) is running.
# I am a Function.

logger实际上是对原有装饰器decorator的一个函数封装,并返回一个装饰器。我们可以将它理解为一个含有参数的闭包。当我们使用`@logger(level=”INFO”)调用的时候,Python 能够发现这一层的封装,并把参数传递到装饰器的环境

用装饰器注册插件

装饰器除了装饰函数,还可以简单地注册一个现有函数并返回未包装的函数

例如,可以使用它来创建轻量级的插件体系结构:

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 pprint
import random
PLUGINS = dict()


def register(func):
"""Register a function as a plug-in"""
PLUGINS[func.__name__] = func
return func


@register
def say_hello(name):
return f"Hello {name}"


@register
def be_awesome(name):
return f"Yo {name}, together we are the awesomest!"


def randomly_greet(name):
greeter, greeter_func = random.choice(list(PLUGINS.items()))
print(f"Using {greeter!r}")
return greeter_func(name)


if __name__ == '__main__':
pprint.pprint(PLUGINS)
print('=' * 30)
randomly_greet("Alice")

可以看到如下输出

1
2
3
4
# {'be_awesome': <function be_awesome at 0x000002923CAC59D8>,
# 'say_hello': <function say_hello at 0x000002923CAC1678>}
# ==============================
# Using 'say_hello'

这样的轻量级插件体系结构的主要好处,就是你不需要去维护插件列表,这个插件列表在插件自动注册时就有了,你只需要用装饰器装饰函数就能添加一个新插件

我们可以用函数globals()查看全局变量,可以发现新产生的插件就在其中

1
2
3
4
5
pprint.pprint(globals())
# # ...Lots of variables not shown here...
# 'randomly_greet': <function randomly_greet at 0x000002923CACD048>,
# 'register': <function register at 0x000002923CAD00D8>,
# 'say_hello': <function say_hello at 0x000002923CAC1678>}

装饰类

装饰不仅可以装饰函数,还可以装饰类

装饰类的函数

装饰类的方式有两种,一种是装饰类中的函数,上面的例子实现的差不多就是这样的方式,还有我们平常常用的几个装饰器,如@classmethod, @staticmethod, @property,其中,前两者不多做介绍,@property装饰器提两点

@property装饰器

@property一般用来实现类内属性的gettersetter方法,这里写一个简单的示例,点到为止

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

@property
def val(self):
return self._val

@val.setter
def secret(self, val):
if val > 0:
self._val = val
else:
raise ValueError('must be a positive integer')

装饰整个类

装饰器用来装饰类和用来装饰函数其实很相似,只不过装饰类时传入装饰器的参数是类

可以构造这样一个类装饰器:singleton–单例模式,该装饰器让被装饰的类只能有一个实例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from functools import wraps

def singleton(cls):
@wraps(cls)
def wrapper(*args, **kwargs):
if not wrapper.instance:
wrapper.instance = cls(*args, **kwargs)
return wrapper.instance
wrapper.instance = None
return wrapper

@singleton
class Admin:
def __init__(self, name):
self.name = name

if __name__ == '__main__':
admin_1 = Admin('Admin1')
admin_2 = Admin('Admin2')
print(admin_1.name) # Admin1
print(admin_2.name) # Admin1
print(id(admin_1) == id(admin_2)) # True

可以看出两个类是完全相同的,特别是通过id(admin_1) == id(admin_2)更能证明这一点

其实我们不难发现,装饰类和装饰函数的装饰器模板大同小异,确实没什么差别,因此这里不再举更多例子

装饰器类

使用装饰器类主要依靠类的__call__方法,当使用 @ 形式将装饰器附加到函数上时,就会调用此方法

示例1:记录函数调用次数

实现一个函数的计数器,记录函数的执行次数,并且每次打印消息,可以装饰在一个计算斐波那契数列的函数上

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

class CountCalls:
def __init__(self, func):
functools.update_wrapper(self, func)
self.func = func
self.num_calls = 0

def __call__(self, *args, **kwargs):
self.num_calls += 1
print(f"Call {self.num_calls} of {self.func.__name__!r}")
return self.func(*args, **kwargs)

@CountCalls
def fibonacci(n):
if n < 2:
return n
return fibonacci(n - 1) + fibonacci(n - 2)

if __name__ == '__main__':
print(fibonacci(6))
# Call 1 of 'fibonacci'
# Call 2 of 'fibonacci'
# ......
# Call 25 of 'fibonacci'
# 8

示例2:记录日志

如果装饰器类带参数,写法还是和函数一样

这里实现一个记录日志的类Logger,每次执行函数就把执行的函数名和执行的时间记录到日志log.txt

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
from functools import wraps
from time import asctime, sleep

class Logger(object):
def __init__(self, log_file='log.txt'):
self.log_file = log_file

def __call__(self, func):
@wraps(func)
def wrapped_function(*args, **kwargs):
log_string = func.__name__ + " was called in " + asctime()
# 打开日志文件并写入
with open(self.log_file, 'a') as f:
f.write(log_string + '\n')
return func(*args, **kwargs)

return wrapped_function

@Logger('log.txt')
def my_func():
print('I am a function.')

if __name__ == "__main__":
my_func()
sleep(5)
my_func()

打开log.txt文件,可以看到如下记录

1
2
my_func was called in Sun Nov  8 10:43:13 2020
my_func was called in Sun Nov 8 10:43:18 2020

使用偏函数与类实现装饰器

一开始提到了,装饰器必须是一个可被调用(callable)的对象,其实偏函数也是 callable 对象

这里介绍一个使用类和偏函数结合实现的装饰器,实现函数延时作用的功能

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

class DelayFunc:
def __init__(self, duration, func):
self.duration = duration
self.func = func

def __call__(self, *args, **kwargs):
print('Please wait for {} seconds...'.format(self.duration))
time.sleep(self.duration)
return self.func(*args, **kwargs)

def eager_call(self, *args, **kwargs):
print('Call without delay')
return self.func(*args, **kwargs)

def delay(duration):
return functools.partial(DelayFunc, duration)

@delay(duration=2)
def add(a, b):
return a + b


if __name__ == '__main__':
print(add(1, 2))
# Please wait for 2 seconds...
# 3

参考文章: