文章目录
前言
大部分的文章在介绍装饰器,基本是以函数实现装饰器为主。事实上,装饰器也可以使用其他方式去实现。
正文
本文中将介绍如何实现装饰器,包括使用类、wrapt模块的方式,并且提醒一些常见错误。
1. 闭包
在函数内部再定义一个函数,并且这个函数用到了外边函数的变量,那么将这个函数以及用到的一些变量称之为闭包.
def func(num):
def inner_func(inner_num):
print(f"in inner_func, inner_num is {inner_num}")
return num + inner_num
return inner_func
res = func(20)
print(res)
print(res(100))
输出:
<function func.<locals>.inner_func at 0x00000139E82B8318>
in inner_func, inner_num is 100
120
使用闭包可以代码复用提高效率,但是由于闭包引用了外部函数的局部变量,因此外部函数的局部变量并不会立刻释放,这会带来额外的开销。
2. 装饰器
2.1 无参装饰器
from time import ctime, sleep
def timer(func):
def wrapper():
print(f"{func.__name__} called at {ctime()}")
func()
return wrapper
@timer
def fn():
print("hello, world")
fn()
sleep(2)
fn()
输出:
fn called at Thu Mar 19 15:56:20 2020
hello, world
fn called at Thu Mar 19 15:56:22 2020 # 间隔2s后再次调用函数
hello, world
2.2 被装饰的函数有参数:
from time import ctime, sleep
from functools import reduce
def timer(func):
def wrapper(*args, **kwargs):
print(f"{func.__name__} called at {ctime()}")
print("传入的参数:", *args, **kwargs)
func(*args, **kwargs)
return wrapper
@timer
def fn(a, b):
print(a + b)
fn(3, 5)
sleep(2)
print("==================================")
fn(2, 4)
输出:
fn called at Thu Mar 19 16:39:40 2020
传入的参数: 3 5
8
==================================
fn called at Thu Mar 19 16:39:42 2020
传入的参数: 2 4
6
2.3 装饰器中的return:
from time import ctime, sleep
def timer(func):
def wrapper():
print(f"{func.__name__} called at {ctime()}")
func()
return wrapper
@timer
def fn():
print("I am fn")
@timer
def get_info():
return '----hahah---'
fn()
print("=============================")
sleep(2)
fn()
print("=============================")
print(get_info())
输出:
I am fn
=============================
fn called at Thu Mar 19 16:41:45 2020
I am fn
=============================
get_info called at Thu Mar 19 16:41:45 2020
None
如果将装饰器中修改为return func()
,即可以打印出get_info()
的return值,因此一般情况下,为了让装饰器更通用,可以加上return。
2.4 带参数的装饰器
from time import ctime, sleep
def timer(pre="hello"):
def decorator(func):
def wrapper():
print(f"{func.__name__} called at {ctime()} {pre}")
return func()
return wrapper
return decorator
@timer()
def fn1():
print("I am fn1")
@timer("python")
def fn2():
print("I am fn2")
@timer("world")
def fn3():
print("I am fn3")
fn1()
print("==============================")
sleep(2)
fn2()
print("==============================")
sleep(2)
fn3()
输出:
fn1 called at Thu Mar 19 16:50:39 2020 hello
I am fn1
==============================
fn2 called at Thu Mar 19 16:50:41 2020 python
I am fn2
==============================
fn3 called at Thu Mar 19 16:50:43 2020 world
I am fn3
3. 类装饰器
绝大多数装饰器的实现都是基于函数,但是这并非是唯一的实现方式。只要是一个callable对象,就可以作为装饰器。
def func():
pass
print(type(func)) # <class 'function'>
print(callable(func)) # True
显然,函数是callable,那么我们也可以通过__call__
这个魔法方法是类变成callable。
class Test(object):
def __call__(self, *args, **kwargs):
print("实现了__call__")
test = Test()
print(type(test)) # <class '__main__.Test'>
print(callable(test)) # True
# 调用test
test() # 实现了__call__
基于这个特性,我们便可以通过类实现装饰器。
下面展示一个类装饰器,用于记录函数的调用次数。
class Counter:
count = 0
def __init__(self, func):
self.func = func
def __call__(self, *args, **kwargs):
Counter.count += 1
return self.func(*args, **kwargs)
@Counter
def fn():
pass
for i in range(10):
fn()
print(fn.count) # 10
相较于纯函数的调用,类装饰器有以下优势:
- 实现有状态的装饰器时,操作类属性比操作闭包内变量更符合直觉、不易出错
- 实现为函数扩充接口的装饰器时,使用类包装函数,比直接为函数对象追加属性更易于维护
- 更容易实现一个同时兼容装饰器与上下文管理器协议的对象(参考 unitest.mock.patch)
当然我们也可以不使用__call__
方法去实现类装饰器,下面我们尝试实现一个简单的@property,使得类中的函数以属性的方式调用。
class Decrator(object):
def __init__(self, func):
self.func = func
def __get__(self, instance, owner):
"""
:param instance: 代表实例,sum中的self
:param owner: 代表类本身,Test类
"""
pass
return self.func(instance)
class Test(object):
def __init__(self):
self.result = 0
@Decrator
def sum(self):
print("hello, world")
t = Test()
t.sum # hello, world
4. wrapt模块
在写装饰器时,往往我们需要面对函数的多层嵌套,可读性不是很高。而且,当我们把原先应用在函数的装饰器,应用到类方法时,会发生错误。
在下面这个例子中,实现了一个生成随机数并注入为函数参数的装饰器。
import random
def provide_number(min_num, max_num):
"""
随机生成一个在 [min_num, max_num] 范围的整数,追加为函数的第一个位置参数
"""
def wrapper(func):
def decorated(*args, **kwargs):
num = random.randint(min_num, max_num)
# 将 num 作为第一个参数追加后调用函数
return func(num, *args, **kwargs)
return decorated
return wrapper
@provide_number(1, 100)
def print_random_number(num):
print(num)
print_random_number() # 随机生成一个1-100范围的整数
然而,当我们把@provide_number
应用到类方法时,得到结果却不是我们想要的。
class Test(object):
@provide_number(1, 100)
def print_random_number(self, num):
print(num)
t = Test()
t.print_random_number()
当调用类方法时,输出的结果是类实例:<__main__.Test object at 0x0000015DCB537E08>
。显然,这不符合我们的要求。
wrapt
模块便可以解决这个问题。使用wrapt你只需要定义一个装饰器函数,但是函数签名是固定的,必须是(wrapped, instance, args, kwargs)
import random
import wrapt
def provide_number(min_num, max_num):
"""
随机生成一个在 [min_num, max_num] 范围的整数,追加为函数的第一个位置参数
"""
@wrapt.decorator
def wrapper(wrapped, instance, args, kwargs):
"""
:param wrapped: 被装饰的函数或类方法
:param instance:
- 如果被装饰者为普通类方法,该值为类实例
- 如果被装饰者为 classmethod 类方法,该值为类
- 如果被装饰者为类/函数/静态方法,该值为None
:param args: 调用时的位置参数
:param kwargs: 调用时的关键字参数
"""
num = random.randint(min_num, max_num)
args = (num, ) + args
return wrapped(*args, **kwargs)
return wrapper
class Test(object):
@provide_number(1, 100)
def print_random_number(self, num):
print(num)
t = Test()
t.print_random_number()
常见错误
1. 使用functools.wraps()装饰内层函数
def decorator(func):
def wrapper(*args, **kwargs):
"""wrapper doc
"""
print(f"func.__name__是:{func.__name__}")
return func(*args, **kwargs)
return wrapper
@decorator
def fn():
"""fn doc
"""
pass
fn() # func.__name__是:fn
print(f"fn.__name__是:{fn.__name__}") # fn.__name__是:wrapper
print(f"fn.__doc__是:{fn.__doc__}") # fn.__doc__是:wrapper doc
使用装饰器后,被装饰函数的函数签名发生了改变,全部变成了内层函数wrapper
的值。这虽然是个小bug,但是有可能会在某些应用场景出问题。此时,我们可以通过functools模块中的wraps解决这个问题。
import functools
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
"""wrapper doc"""
print(f"func.__name__是:{func.__name__}")
return func(*args, **kwargs)
return wrapper
@decorator
def fn():
"""fn doc"""
pass
fn() # func.__name__是:fn
print(f"fn.__name__是:{fn.__name__}") # fn.__name__是:fn
print(f"fn.__doc__是:{fn.__doc__}") # fn.__doc__是:fn doc
这样处理后,装饰器就不会影响被装饰函数了。
2. 修改外层变量时记得使用 nonlocal
装饰器是对函数对象的一个高级应用。在编写装饰器的过程中,你会经常碰到内层函数需要修改外层函数变量的情况。就像下面这个装饰器一样:
import functools
def counter(func):
"""装饰器:记录并打印调用次数"""
count = 0
@functools.wraps(func)
def decorated(*args, **kwargs):
# 次数累加
count += 1
print(f"Count: {count}")
return func(*args, **kwargs)
return decorated
@counter
def fn():
pass
fn()
为了统计函数调用次数,我们需要在decorated函数内部修改外层函数定义的count变量的值。但是,上面这段代码是有问题的,在执行它时解释器会报错:
Traceback (most recent call last):
File "test.py", line 21, in <module>
fn()
File "test.py", line 11, in decorated
count += 1
UnboundLocalError: local variable 'count' referenced before assignment
这个错误是由 counter 与 decorated 函数互相嵌套的作用域引起的。当解释器执行到 count += 1 时,并不知道 count 是一个在外层作用域定义的变量,它把 count 当做一个局部变量,并在当前作用域内查找。最终却没有找到有关 count 变量的任何定义,然后抛出错误。
为了解决这个问题,我们可以使用nonlocal
关键字。
@functools.wraps(func)
def decorated(*args, **kwargs):
# 次数累加
nonlocal count
count += 1
print(f"Count: {count}")
return func(*args, **kwargs)