编写高质量Python (第26条) 用 functools.wraps 定义函数装饰器

第26条 用 functools.wraps 定义函数装饰器

​ Python 中有一个特殊写法，可以用装饰器来封装某个函数，从而让函数在执行这个函数之前与执行完这个函数之后，分别运行某些代码。这意味着，调用者传给参数的参数值、函数返回给调用者的值，以及函数抛出的异常，都可以由装饰器访问并修改。这是个很有用 的机制，能够保证用户以正确的方式使用函数，也能用来调试程序或实现函数注册功能，此外还有很多用途。

​ 例如，假设我们要把函数执行时收到的参数与返回的值记录下来。这在调试递归函数是很有用的，因为我们要知道，这个函数执行每一层递归时，输入的是什么参数，返回的是什么值。下面我们就定义这样一个装饰器，在实现这个修饰器时，用 *args 与 **kwargs 表示受修饰的原函数 func 所收到的参数（参见 第22条和第23条）。

def trace(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        print(f'{func.__name__}({args!r},{kwargs!r})'
              f' -> {result!r}')
        return result

    return wrapper

​ 写好之后，我们用 @ 符号把修饰器运用在想要调试的函数上面。

@trace
def fibonacci(n):
    """Return the n-th Fibonacci number """
    if n in (0, 1):
        return n
    return (fibonacci(n - 2) + fibonacci(n - 1))

​ 这样写，相当于先把受修饰的函数传给修饰器，然后将修饰器所返回的值赋给原来那个函数。这样的话，如果我们继续通过原来的那个名字调用函数，那么执行的就是修饰之后的函数。

fibonacci = trace(fibonacci)

​ 修饰过的 fibonacci 函数，会在执行自身的代码之前，先执行 wrapper 里位于 func(*args, **kwargs) 那一行之前的逻辑；并且在执行完自身的代码后，执行 wrapper 里位于 func (*args, **kwargs) 那一行之后的逻辑。本例中，它会在执行完自身的代码之后，打印这次执行所用的参数与返回值，这样就能看到整个递归栈的情况了。

fibonacci(4)

>>>
fibonacci((0,),{}) -> 0
fibonacci((1,),{}) -> 1
fibonacci((2,),{}) -> 1
fibonacci((1,),{}) -> 1
fibonacci((0,),{}) -> 0
fibonacci((1,),{}) -> 1
fibonacci((2,),{}) -> 1
fibonacci((3,),{}) -> 2
fibonacci((4,),{}) -> 3

​ 这样写确实能满足需求，但是会带来一个我们不愿意看到的副作用。修饰器返回的那个值，也就是刚才调用的 fibonacci ，它的名字并不叫 “fibonacci”。

print(fionacci)

>>>
<function trace.<locals>.wrapper at 0x1090ad120>

​ 这种现象解释起来并不困难。trace 函数返回的，是它里面定义的 wrapper 函数，所以，当我们把这个返回值赋给 fibonacci 之后，fibonacci 这个名称所表示的自然就是 wrapper 了。问题在于，这样可能会干扰那些想要利用 introspection 机制来运作的工具，例如调试器（debugger）（参见 第80条）。

​ 例如，如果用内置的 help 函数来查看修饰后的 fibonacci，那么打印出来的并不是我们想看的帮助文档，他本来该打印前面定义时写的那行 'Return the n-th Fibonacci number’文本才对。

help(fabonacci)

>>>
Help on function wrapper in module __main__:

wrapper(*args, **kwargs)

​ 对象序列化器(object serializer, 参见 第68条) 也无法正常运作，因为它不能确定受修饰函数的位置。

import pickle
pickle.dumps(fibonacci)

>>>
Traceback ...
AttributeError: Can't pickle local object 'trace.<locals>.wrapper'

​ 要像解决这些问题，可以改用 functools 内置模块之中的 wraps 辅助函数来实现。wraps 本身也是个修饰器，它可以帮助你编写自己的装饰器。把它运作到 wrapper 函数上面，它就会将重要的元数据 (metadata) 全都从内部函数复制到外部函数。

from functools import wraps


def trace(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        print(f'{func.__name__}({args!r},{kwargs!r})'
              f' -> {result!r}')
        return result

    return wrapper


@trace
def fibonacci(n):
    """Return the n-th Fibonacci number """
    if n in (0, 1):
        return n
    return (fibonacci(n - 2) + fibonacci(n - 1))

​ 现在我们就可以通过 help 函数看到正确的文档了。虽然原来的 fibonacci 函数现在封装在装饰器上，但我们还可以看到它的文档。

help(fibonacci)

>>>
Help on function fibonacci in module __main__:

fibonacci(n)
    Return the n-th Fibonacci number

​ 对象序列化，也正常了。

print(pickle.dumps(fibonacci))

>>>
b'\x80\x04\x95\x1a\x00\x00\x00\x00\x00\x00\x00\x8c\x08__main__\x94\x8c\tfibonacci\x94\x93\x94.'

​ 除了这里讲到的几个方面之外，Python 函数还有很多属性（例如 __name__,_moudle_,__annotations__)也应该在受到封装时得以保留，这样才能让相关的接口正常运作。wraps 可以帮助保存这些属性，使程序表现出正确的行为。