我后悔不早知道的8个Python装饰器技巧

1) @ 语法等价性

@add_exclamation
def greet(name):
    return 'hello' + name

• add_exclamation 是装饰器函数。
• greet 是被装饰的函数（我们叫它“被装饰者”）。

我们用 @ 符号来装饰一个函数，使用装饰器函数。这种 @ 语法和下面这段代码完全一样：

def greet(name):
    return 'hello' + name
    
greet = add_exclamation(greet)

这就是我们在使用装饰器函数 add_exclamation 装饰 greet 函数时背后发生的事情。

2) 我们为什么使用装饰器

装饰器函数本身就是函数——它可以修改被装饰者的功能，而无需改变被装饰者的源代码。

@add_exclamation
def greet(name):
    return 'hello' + name

再次强调，add_exclamation 是装饰器函数，greet 是被装饰者。

• add_exclamation 修改了 greet 的行为。
• 它会在 greet 的返回值后面加上一个感叹号。
• 经过装饰后，greet('tom') 将会返回 'hello tom!'。

这里要注意的是，我们实际上并没有在 greet 函数内部添加感叹号。

这一点很重要，因为装饰器的目的就是装饰多个函数，并为每个被装饰的函数增加额外的功能。

3) 理解装饰器的简便方法

再次强调，使用 @ 语法和使用 _func = decorator(func) 是一样的。

这意味着我们的装饰器函数 add_exclamation 接受一个函数作为参数，并返回一个修改过的函数。

明白了这一点，让我们尝试实现 add_exclamation 装饰器函数，它会在被装饰者返回值的后面加一个感叹号。

def add_exclamation(func):
    def wrapper(name):
        return func(name) + '!'
    return wrapper

1. greet = add_exclamation(greet)
2. 变量 _func_ 被赋值为我们函数 greet。
3. 创建了一个内嵌函数 _wrapper_。_wrapper_ 接收和 greet 相同的输入。
4. _wrapper_ 返回 greet 的结果，但多了一个感叹号。
5. _wrapper_ 本身是由函数调用 add_exclamation(greet) 返回的。
6. 变量 greet 被赋值为这个新的返回函数。

这样一来，greet 的返回值就多了一个感叹号，而我们并不需要手动修改 greet 的源代码。

@add_exclamation
def greet(name):
    return 'hello' + name

print(greet('tom')) # hello tom!

4) 具有两个内嵌函数的高级装饰器

但是如果我们要添加不同的符号，比如问号、逗号等等呢？我们需要为每一个符号创建一个装饰器函数吗？

好消息是：不用。我们可以把装饰器设计得更加通用：

@add_symbol('!')
def greet(name):
    return 'hello' + name

print(greet('tom')) # hello tom!

同样地，这里的 @ 语法等同于：

greet = add_symbol('!')(greet)

这意味着我们的 add_symbol 装饰器函数需要有两个内嵌函数。现在我们来实现这个装饰器函数：

def add_symbol(symbol):
    def decorator(func):
        def wrapper(name):
            return func(name) + symbol
        return wrapper
    return decorator

5) 类也可以是装饰器

如果你不想处理包含两个内嵌函数的复杂装饰器，还有一种方法——我们可以用类作为装饰器。

要做到这一点，我们需要实现 __call__ 魔法方法，它定义了当我们像调用函数那样调用一个对象时的行为。

class add_symbol:
    def __init__(self, symbol):
        self.symbol = symbol

    def __call__(self, func):
        def wrapper(name):
            return func(name) + self.symbol
        return wrapper
        
@add_symbol('!')
def greet(name):
    return 'hello' + name

print(greet('tom')) # hello tom!

这里，我们用类重写了之前的例子，而不是用带有两个内嵌函数的函数。

同样地，使用 @ 符号等同于这样操作：

greet = add_symbol('!')(greet)

• 第一组括号初始化我们的对象并调用 __init__。
• 第二组括号调用我们的对象并执行 __call__，它返回一个像普通装饰器那样的内嵌 _wrapper_ 函数。

6) functools.wraps

假设我们有一个简单的装饰器，它会在函数的返回值后面加一个感叹号。但是当我们打印 greet 的函数元数据时，会得到这样的结果：

@add_symbol('!')
def greet(name):
    """ 向某人打招呼 """
    return 'hello' + name

print(greet('tom')) # hello tom!


print(greet.__name__) # wrapper
print(greet.__doc__) # None

这是因为，在 greet = add_exclamation(greet) 中，

• add_exclamation(greet) 返回其内嵌函数 _wrapper_。
• greet 被赋值为内嵌函数 _wrapper_。
• 因此，greet 失去了它的函数元数据，如名称和文档字符串。

我们可以使用内置的 functools.wraps 来保留这些函数元数据。

from functools import wraps

def add_exclamation(func):
    @wraps(func)
    def wrapper(name):
        return func(name) + '!'
    return wrapper

注意——唯一的变化是我们用 @wraps(func) 装饰了 _wrapper_。这个变化迫使 _func_ 的元数据传递给 _wrapper_。

现在，我们的元数据被保留下来了。

print(greet.__name__) # greet
print(greet.__doc__) # 向某人打招呼

7) 使用 __wrapped__ 来取消装饰

当我们装饰一个函数后，该函数就会永远改变，我们似乎失去了对原始未装饰函数的访问。

这里，greet 的返回值将永久多出一个感叹号。

但如果我们要获取没有感叹号的原始 greet 函数，我们就没辙了。除非……

1. 我们使用 functools.wraps 来装饰装饰器中的 _wrapper_ 函数。
2. 我们使用 .__wrapped__ 来访问原始未装饰的函数。

print(greet('tom')) # hello tom!
print(greet.__wrapped__('tom')) # hello tom

这在单元测试中很有用，因为我们可能需要提取原始未装饰的函数。

8) 用来装饰类的装饰器

装饰器既可以用来装饰函数，也可以用来装饰类。当然，用于装饰类的装饰器与装饰函数的装饰器会有很大不同。

假设我们想要创建一个类装饰器，它会给被装饰的每个类添加一个共同的方法 _hello_。

def add_hello(cls):
    cls.hello = lambda self: print('hello')
    return cls

@add_hello
class Dog: pass

Dog().hello() # hello

在装饰器函数 add_hello 中，我们简单地向类分配了一个方法，并返回该类本身。

如果我们用 add_hello 装饰多个类，那么这些类都会默认包含 _hello_ 方法。

结论

希望今天你至少学到了一些关于 Python 装饰器的新知识。