【Python装饰器】functools.wraps函数保留被装饰函数的元信息

前言

装饰器一般被用于修饰函数，为被修饰的函数增添某些功能，其输入一般为函数，输出为同一个函数，或者另一不同的函数。除注册装饰器外，大多数装饰器会返回与被装饰函数不同的函数对象。另一方面，由于装饰器内部定义中会返回函数，因此涉及到嵌套函数定义，“闭包”是嵌套函数能够正确运行的基础之一，闭包可以简单理解为“内部嵌套的函数的作用域有所外延，可以引用闭包中的自有变量”。

下文中定义几个基本名词概念方便描述，假设装饰器返回的函数对象叫做wrapper，被装饰的对象叫做wrapped。装饰器经常通过语法糖的形式装饰函数：

# ******* 语法糖形式等价于func = decorate(func) ********
@decorate  
def func(x):
	pass

使用这种修饰方式有一个的副作用，就是wrapper对象会丢失掉原有输入参数func函数的属性信息，因为返回的对象wrapper是一个全新的对象（大多数情况下），有自己独立的属性信息。那么如何为函数添加额外的功能，同时又保留函数的元信息，解决方案——使用wraps函数修饰函数。要想理解wraps函数作为装饰器的用法，需要先了解functools中的update_wrapper函数与partial类。
 

update_wrapper函数

update_wrapper函数是解决前言问题的核心，字面上update_wrapper函数就是更新wrapper的信息，具体来讲就是将被修饰的函数对象的属性信息，赋值或者更新到wrapper中，wrapper源码定义如下：

def update_wrapper(wrapper,
                   wrapped,
                   assigned = WRAPPER_ASSIGNMENTS,
                   updated = WRAPPER_UPDATES):
    """Update a wrapper function to look like the wrapped function

       wrapper is the function to be updated
       wrapped is the original function
       assigned is a tuple naming the attributes assigned directly
       from the wrapped function to the wrapper function (defaults to
       functools.WRAPPER_ASSIGNMENTS)
       updated is a tuple naming the attributes of the wrapper that
       are updated with the corresponding attribute from the wrapped
       function (defaults to functools.WRAPPER_UPDATES)
    """
    # WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__qualname__', '__doc__', '__annotations__')
    for attr in assigned:
        try:
            value = getattr(wrapped, attr)
        except AttributeError:
            pass
        else:
            setattr(wrapper, attr, value)
   
   # WRAPPER_UPDATES = ('__dict__',)
    for attr in updated:
        getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
        
    # Issue #17482: set __wrapped__ last so we don't inadvertently copy it
    # from the wrapped function when updating __dict__
    wrapper.__wrapped__ = wrapped   # wrapper.__wrapped__属性中保存着wrapped的值，一般就是被修饰函数
     
    # Return the wrapper so this can be used as a decorator via partial()
    return wrapper

wrapper_update的作用是将参数值wrapped的assigned、updated指定的属性信息赋值给wrapper，返回值是更新信息的wrapper。试想前言中提到的问题，当wrapped是func时，那么wrapper将保留func的信息，因此展示的就是func的信息，实例如下：

from operator import add
from functools import update_wrapper, partial, wraps

# 首先定义一个装饰器
def my_decorator(func):
    """ add say utility to builtins add """
    print('begin decorator!')
    def wrapper(*args, **kwargs):
        """ i'm function inner my_decorator """
        print('say hello!')
        result = func(*args, **kwargs)
        print('say bye!')
        return result
    print('end decorator!')
    return wrapper

my_add = my_decorator(add)

下图给了一些关于add与my_add的信息展示：

首先可以看到打印的“begin decorate”与“end decorator”信息，说明my_decorator是正常运行的，然后展示了my_add的名称与文档字符串信息，与my_decorator定义体保持一致，后续展示了内置的add函数名称与文档字符串信息。
 

下图展示装饰器赋予add内置函数的新功能，say hello与say bye。

现在使用update_wrapper函数将add函数的属性信息复制到my_add，assigned与updated关键字参数用默认设置就可以了：

可以发现my_add1现在不仅有了“say”功能，而且my_add1的属性信息就是内置函数add的属性信息。
 

partial类

首先注意partial是一个类，截取后的部分源码如下：

class partial:
    """New function with partial application of the given arguments
    and keywords.
    """
    __slots__ = "func", "args", "keywords", "__dict__", "__weakref__"

    def __new__(cls, func, /, *args, **keywords):
        if not callable(func):
            raise TypeError("the first argument must be callable")

        if hasattr(func, "func"):
            args = func.args + args
            keywords = {**func.keywords, **keywords}
            func = func.func

        self = super(partial, cls).__new__(cls)

        self.func = func
        self.args = args
        self.keywords = keywords
        return self

    def __call__(self, /, *args, **keywords):
        keywords = {**self.keywords, **keywords}
        return self.func(*self.args, *args, **keywords)

由于定义了__call__方法，partial实例就是一个可调用对象，并且__call__中实际就是以初始化的参数值调用实例化的func函数对象。对于partial实例的调用可以这样理解，本来调用func函数需要提供多个参数值，但在实例化partial对象时，提供了一些位置参数与关键字参数，这些参数值就被固定在patial实例对象里。当调用partial对象时，就只需要提供少量参数就可以了，可以类比高数中的求偏导的概念。下面给出partial的应用实例：

add_fix = partial(add, 99)  # 固定加法运算中的一个参数
print(add_fix(1))  # 返回100， == 99+1
print(add_fix(2))  # 返回101， == 99+2

 

wraps函数

wraps函数就是结合partial类与update_wrapper的函数。wraps的源码定义如下：

def wraps(wrapped,
          assigned = WRAPPER_ASSIGNMENTS,
          updated = WRAPPER_UPDATES):
    """Decorator factory to apply update_wrapper() to a wrapper function

       Returns a decorator that invokes update_wrapper() with the decorated
       function as the wrapper argument and the arguments to wraps() as the
       remaining arguments. Default arguments are as for update_wrapper().
       This is a convenience function to simplify applying partial() to
       update_wrapper().
    """
    # 注意：update_wrapper是上面介绍的update_wrapper函数
    return partial(update_wrapper, wrapped=wrapped, assigned=assigned, updated=updated)

partial实例化过程与update_wrapper函数调用非常类似，就是少了个wrapper参数。如果调用partial实例，并提供wrapper参数，则相当于update_wrapper(wrapper, wrapped=wrapped, assigned=assigned, updated=updated)， 后面的三个关键字参数值都是在partial实例化时存储的。

注意wraps需要放在wrapper之前，如下所示：

@wraps(func)会依次执行wraps(func) -> partial(func, *args, **kwargs)， 返回的partial实例是可调用的， 假设为p，p(wrapper)等价于 update_wrapper(wrapper, **kwargs)，红框中wraps(func)返回下列结果：

结果是一个partial实例，实际上就是固定了wrapped、assigned、updated三个关键字参数update_wrapper。当用返回的partial实例再去修饰wrapper时，相当于wrapper = partial(wrapper), 就将wrapped的属性信息复制、更新到wrapper属性信息中。因此可以看到加了wraps的my_add对象，其名字等属性就是内置函数add的属性，这样就解决了前言中的问题。注意“wraps本质是装饰器工程函数，因为其返回值是partial实例，partial实例可以作为装饰器修饰其它函数”。
 

带参数的装饰器

装饰器的输入参数一般是函数，如果需要定义带参数的装饰器，该如何实现？——在装饰器定义的外层定义一个“装饰器工厂函数”，所需的装饰器的参数放置在“装饰器工程函数”中，内层的装饰器可以访问外层“装饰器工程函数”的自由变量。如下例所示：

def decorator_factory(count=True):
    """ 定义装饰器工厂函数 """
    def decorator(func):
        """ 真正的装饰器 """
        @functools.wraps(func)
        def wrapper(*args):
            """ 装饰器内置函数 """
            if count: # 引用自由变量
                start = time.time()
                res = func(*args)
                print('total cost time :{}s'.format(time.time() - start))
                return res
            else:
                return func(*args)
        return wrapper
    return decorator

接下来以调用的方式装饰函数：

@decorator_factory(count=True)
def f1(n):
    if n < 2:
        return n
    return f1(n-1) + f1(n-2)

当count=True时，表示启用计时功能，如下所示：

当count=False时，则不启用计时功能，相当于f函数：

总之实现带参数的装饰器最重要的有两点，1）定义装饰器工厂函数，将参数放入工厂函数定义中，工厂函数的返回值为实际需要的装饰器对象；2）使用语法糖形式时，一定要以函数调用的方式调用装饰器工厂函数，因为是调用，所以工厂函数会返回装饰器，然后再用返回的装饰器修饰目标函数。

参考资料

1. 流程的python——第7章函数装饰器与闭包
2. https://docs.python.org/3.7/librry/functools.html