functools.wraps用法

介绍

在Python中，functools.wraps()是一个装饰器，它的主要作用是将被装饰函数的一些属性值（如__name__、__doc__等）赋值给装饰器中的函数。这个功能主要用于在使用装饰器时保持函数的元数据。

示例

让我们通过一个例子来看看它是如何工作的：

def my_decorator(f):
    def wrapper(*args, **kwds):
        print('Calling decorated function...')
        return f(*args, **kwds)
    return wrapper

@my_decorator
def example():
    """Docstring for example function"""
    print('Called example function')

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

在这个例子中，我们定义了一个装饰器my_decorator，它在调用原函数前打印一条消息。然后我们用这个装饰器装饰了函数example。

然后我们尝试打印出example函数的名字和文档字符串。但是，运行这段代码，你会发现打印出的函数名字是"wrapper"，而不是"example"，而且example函数的文档字符串也丢失了。

这就是因为当我们使用装饰器时，实际上是定义了一个新的函数（在本例中就是wrapper函数）来代替原来的函数。新函数的名字和文档字符串与原函数不同，所以当我们试图访问这些属性时，得到的是新函数的属性值。

为了解决这个问题，我们可以使用functools.wraps()来修复它：

from functools import wraps

def my_decorator(f):
    @wraps(f)
    def wrapper(*args, **kwds):
        print('Calling decorated function...')
        return f(*args, **kwds)
    return wrapper

@my_decorator
def example():
    """Docstring for example function"""
    print('Called example function')

print(example.__name__) # example
print(example.__doc__) # Docstring for example function

现在，打印出的函数名字是"example"，文档字符串也恢复了。这是因为@wraps(f)将原函数的元数据复制到了装饰器函数中。

所以，functools.wraps()的实现原理就是在装饰器内部将原函数的元数据复制到新函数中，从而让新函数看起来更像原函数。

源码剖析

wraps中返回了update_wrapper函数：

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().
    """
    return partial(update_wrapper, wrapped=wrapped,
                   assigned=assigned, updated=updated)

update_wrapper调用时将wrapped所有的属性赋值给wrapper（其中wrapper是装饰器中的函数，wrapped是被装饰的函数）：

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)
    """
    for attr in assigned:
        try:
            value = getattr(wrapped, attr)
        except AttributeError:
            pass
        else:
            setattr(wrapper, attr, value)
    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
    # Return the wrapper so this can be used as a decorator via partial()
    return wrapper

参考资料：python@wraps实现原理