python@wraps实现原理

Posted 2021-01-29 黄土地上的黑石头

@wraps作用

python中的装饰器装饰过的函数其实就不是函数本身了，我们可以看看下面的例子

import time
def timmer(func):
    """timmer doc"""
    def inner(*args, **kwargs):
        """inner doc"""
        start = time.time()
        res = func(*args, **kwargs)
        end = time.time()
        print("函数运行时间为 %s" % (end - start))
        return res
    return inner


@timmer
def func_test():
    """func_test doc"""
    time.sleep(2)
    return

print(func_test.__name__)  # inner
print(func_test.__doc__)  # inner doc

按我们正常的思维，func_test.__name__应该拿到的就是“func_test”，所以这个结果就印证了上面的第一句话，但是这是我们加一个@wraps，就会发现好像一切都正常了：

import time
from functools import wraps
def timmer(func):
    """timmer doc"""
    @wraps(func)
    def inner(*args, **kwargs):
        """inner doc"""
        start = time.time()
        res = func(*args, **kwargs)
        end = time.time()
        print("函数运行时间为 %s" % (end - start))
        return res
    return inner


@timmer
def func_test():
    """func_test doc"""
    time.sleep(2)
    return

print(func_test.__name__)  # func_test
print(func_test.__doc__)  # func_test doc

@wraps的实现原理

为了方便理解，我把源码和例子放在了一起，这样的话我们看着会方便：

import time
from functools import partial

WRAPPER_ASSIGNMENTS = (‘__module__‘, ‘__name__‘, ‘__qualname__‘, ‘__doc__‘,
                       ‘__annotations__‘)
WRAPPER_UPDATES = (‘__dict__‘,)


def update_wrapper(wrapper,  # inner
                   wrapped,  # func_test
                   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)
    """
    print(‘update_wrapper 执行...‘)
    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()
    print(‘update_wrapper 执行结束‘)
    return 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().
    """
    print(‘wraps 执行...‘)
    print(‘wraps 执行结束‘)  # 纯粹为了打印出来的结果好理解
    return partial(update_wrapper, wrapped=wrapped,
                   assigned=assigned, updated=updated)


def timmer(func):
    print(‘timmer 执行...‘)

    @wraps(func)  # inner = update_wrapper的返回值
    def inner(*args, **kwargs):
        start = time.time()
        res = func(*args, **kwargs)
        end = time.time()
        print("函数运行时间为 %s" % (end - start))
        return res

    print(‘timmer 执行结束‘)  # 当然不是真正的结束，执行完下一行才结束
    return inner


@timmer
def func_test():
    print("func_test 执行...")
    time.sleep(2)
    print("func_test 运行结束")
    return


func_test()

"""
打印结果如下：
    timmer 执行...
    wraps 执行...
    wraps 执行结束
    update_wrapper 执行...
    update_wrapper 执行结束
    timmer 执行结束
    func_test 执行...
    func_test 运行结束
    函数运行时间为 2.0000197887420654


从打印的结果我们可以看出，@语法会在函数定义或者说模块初始化阶段(可能称呼不对，以后回来改)就执行了
"""

上面的例子中我加了很多打印，主要是为了提醒一下在func_test()函数执行之前，@语法已经执行了。

其实原理很简单，用了一个偏函数，去执行update_wrapper，真正起作用的也是这个函数，func_test执行之前，update_wrapper函数就会把inner函数的好多属性（示例中WRAPPER_ASSIGNMENTS，WRAPPER_UPDATES指向的属性 ，还有__wrapped__属性）全部其换成func_test的属性。