探究functools模塊wraps裝飾器的用途

《A Byte of Python》17.8節講decorator的時候，用到了functools模塊中的一個裝飾器：wraps。因為之前沒有接觸過這個裝飾器，所以特地研究了一下。

何謂“裝飾器”？

《A Byte of Python》中這樣講：

“Decorators are a shortcut to applying wrapper functions. This is helpful to “wrap” functionality with the same code over and over again.”

《Python參考手冊（第4版）》6.5節描述如下：

“裝飾器是一個函數，其主要用途是包裝另一個函數或類。這種包裝的首要目的是透明地修改或增強被包裝對象的行為。”

Python官方文檔中這樣定義：

“A function returning another function, usually applied as a function transformation using the @wrapper syntax. Common examples for decorators are classmethod() and staticmethod().”

讓我們來看一下《Python參考手冊》上6.5節的一個例子（有些許改動）：

# coding: utf-8
# Filename: decorator_wraps_test.py
# 2014-07-05 18:58
import sys

debug_log = sys.stderr

def trace(func):
    if debug_log:
        def callf(*args, **kwargs):
            """A wrapper function."""
            debug_log.write('Calling function: {}\n'.format(func.__name__))
            res = func(*args, **kwargs)
            debug_log.write('Return value: {}\n'.format(res))
            return res
        return callf
    else:
        return func

@trace
def square(x):
    """Calculate the square of the given number."""
    return x * x

if __name__ == '__main__':
    print(square(3))

輸出：
Calling function: square
Return value: 9
9

這個例子中，我們定義了一個裝飾器trace，用於追蹤函數的調用過程及函數調用的返回值。如果不用裝飾器語法，我們也可以這樣寫：

def _square(x):
    return x * x

square = trace(_square)

上面兩段代碼，使用裝飾器語法的版本和不用裝飾器語法的版本實際上是等效的。只是當我們使用裝飾器時，我們不必再手動調用裝飾器函數。

嗯。trace裝飾器看起來棒極了！假設我們把如上代碼提供給其他程序員使用，他可能會想看一下square函數的幫助文檔：

>>> from decorator_wraps_test import square
>>> help(square) # print(square.__doc__)
Help on function callf in module decorator_wraps_test:

callf(*args, **kwargs)
    A wrapper function.

看到這樣的結果，使用decorator_wraps_test.py模塊的程序員一定會感到困惑。他可能會帶着疑問敲入如下代碼：

>>> print(square.__name__)
callf

這下，他可能會想看一看decorator_wraps_test.py的源碼，找一找問題究竟出現在了哪里。我們知道，Python中所有對象都是“第 一類”的。比如，函數（對象），我們可以把它當作普通的數據對待：我們可以把它存儲到容器中，或者作為另一個函數的返回值。上面的程序中，在 debug_log為真的情況下，trace會返回一個函數對象callf。這個函數對象就是一個“閉包”，因為當我們通過：

def _square(x): return x * x
square = trace(_square)

把trace返回的callf存儲到square時，我們得到的不僅僅是callf函數執行語句，還有其上下文環境:

>>> print('debug_log' in square.__globals__)
True
>>> print('sys' in square.__globals__)
True

因此，使用裝飾器修飾過的函數square，實際上是一個trace函數返回的“閉包”對象callf，這就揭示了上面help(square)以及print(square.__name__)的輸出結果了。

那么，怎樣才能在使用裝飾器的基礎上，還能讓help(square)及print(square.__name__)得到我們期待的結果呢？這就是functools模塊的wraps裝飾器的作用了。

讓我們先看一看效果：

# coding: utf-8
# Filename: decorator_wraps_test.py
# 2014-07-05 18:58
import functools
import sys

debug_log = sys.stderr

def trace(func):
    if debug_log:
        @functools.wraps(func)
        def callf(*args, **kwargs):
            """A wrapper function."""
            debug_log.write('Calling function: {}\n'.format(func.__name__))
            res = func(*args, **kwargs)
            debug_log.write('Return value: {}\n'.format(res))
            return res
        return callf
    else:
        return func

@trace
def square(x):
    """Calculate the square of the given number."""
    return x * x

if __name__ == '__main__':
    print(square(3))
    print(square.__doc__)
    print(square.__name__)

輸出：

Calling function: square
Return value: 9
9
Calculate the square of the given number.
square

很完美！哈哈。這里，我們使用了一個帶參數的wraps裝飾器“裝飾”了嵌套函數callf，得到了預期的效果。那么，wraps的原理是什么呢？

首先，簡要介紹一下帶參數的裝飾器：

>>> def trace(log_level):
    def impl_f(func):
        print(log_level, 'Implementing function: "{}"'.format(func.__name__))
        return func
    return impl_f

>>> @trace('[INFO]')
def print_msg(msg): print(msg)

[INFO] Implementing function: "print_msg"
>>> @trace('[DEBUG]')
def assert_(expr): assert expr

[DEBUG] Implementing function: "assert_"
>>> print_msg('Hello, world!')
Hello, world!

這段代碼定義了一個帶參數的trace裝飾器函數。因此：

@trace('[INFO]')
def print_msg(msg): print(msg)

等價於：

temp = trace('[INFO]')
def _print_msg(msg): print(msg)
print_msg = temp(_print_msg)

相信這樣類比一下，帶參數的裝飾器就很好理解了。（當然，這個例子舉得並不好。《Python參考手冊》上有一個關於帶參數的裝飾器的更好的例子，感興趣的童鞋可以自己看看 。）

接下來，讓我們看看wraps這個裝飾器的代碼吧！

讓我們先找到functools模塊文件的路徑：

>>> import functools
>>> functools.__file__
'D:\\Program Files\\Python34\\lib\\functools.py'

下面，把wraps相關的代碼摘錄出來：

WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__qualname__', '__doc__',
                       '__annotations__')
WRAPPER_UPDATES = ('__dict__',)
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

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)

從代碼中可以看到，wraps是通過functools模塊中另外兩個函數：partial和update_wrapper來實現其功能的。讓我們看一看這兩個函數：

1. partial函數

partial函數實現對函數參數進行部分求值（《Python參考手冊》中4.9有這么一句：函數參數的部分求值與叫做柯里化（currying）的過程關系十分密切。雖然不是太明白，但感覺很厲害的樣子！2014-07-07 15:05追加內容：在百度博客中，zotin大哥回復了我，並對函數式編程中柯里化概念做了一些解釋。）：

>>> from functools import partial
>>> def foo(x, y, z):
    print(locals())
>>> foo(1, 2, 3)
{'z': 3, 'y': 2, 'x': 1}
>>> foo_without_z = partial(foo, z = 100)
>>> foo_without_z
functools.partial(<function foo at 0x00000000033FC6A8>, z=100)
>>> foo_without_z is foo
False
>>> foo_without_z(10, 20)
{'z': 100, 'y': 20, 'x': 10}

這里，我們通過partial為foo提供參數z的值，得到了一個新的“函數對象”（這里之所以加個引號是因為foo_without_z和一般的函數對象有些差別。比如，foo_without_z沒有__name__屬性。）foo_without_z。因此，本例中：

foo_without_z(10, 20)

等價於：

foo(10, 20, z = 100)

（比較有趣的一點是，foo_without_z沒有__name__屬性，而其文檔字符串__doc__也和partial的文檔字符串很相像。此外， 我認為，這里的partial和C++標准庫中的bind1st、bind2nd這些parameter binders有異曲同工之妙。這里沒有把partial函數的實現代碼摘錄出來，有興趣的童鞋可以自己研究一下它的工作原理。）

因此，wraps函數中：

    return partial(update_wrapper, wrapped=wrapped,
                   assigned=assigned, updated=updated)

實際上是返回一個對update_wrapper進行部分求值的“函數對象”。因此，上例中使用了wraps裝飾器的decorator_wraps_test.py的等價版本如下：

def trace(func):
    if debug_log:
        def _callf(*args, **kwargs):
            """A wrapper function."""
            debug_log.write('Calling function: {}\n'.format(func.__name__))
            res = func(*args, **kwargs)
            debug_log.write('Return value: {}\n'.format(res))
            return res

        _temp = functools.wraps(func)
        callf = _temp(_callf)
        return callf
    else:
        return func

對wraps也進行展開：

def trace(func):
    if debug_log:
        def _callf(*args, **kwargs):
            """A wrapper function."""
            debug_log.write('Calling function: {}\n'.format(func.__name__))
            res = func(*args, **kwargs)
            debug_log.write('Return value: {}\n'.format(res))
            return res

        _temp = functools.partial(functools.update_wrapper,
                                  wrapped = func,
                                  assigned = functools.WRAPPER_ASSIGNMENTS,
                                  updated = functools.WRAPPER_UPDATES)
        callf = _temp(_callf)
        return callf
    else:
        return func

最后，對partial的調用也進行展開：

def trace(func):
    if debug_log:
        def _callf(*args, **kwargs):
            """A wrapper function."""
            debug_log.write('Calling function: {}\n'.format(func.__name__))
            res = func(*args, **kwargs)
            debug_log.write('Return value: {}\n'.format(res))
            return res

        callf = functools.update_wrapper(_callf,
                                         wrapped = func,
                                         assigned = functools.WRAPPER_ASSIGNMENTS,
                                         updated = functools.WRAPPER_UPDATES)

        return callf
    else:
        return func

這次，我們看到的是很直觀的函數調用：用_callf和func作為參數調用update_wrapper函數。

2. update_wrapper函數

update_wrapper做的工作很簡單，就是用參數wrapped表示的函數對象（例如：square）的一些屬性（如：__name__、 __doc__）覆蓋參數wrapper表示的函數對象（例如：callf，這里callf只是簡單地調用square函數，因此可以說callf是 square的一個wrapper function）的這些相應屬性。

因此，本例中使用wraps裝飾器“裝飾”過callf后，callf的__doc__、__name__等屬性和trace要“裝飾”的函數square的這些屬性完全一樣。

經過上面的分析，相信你也了解了functools.wraps的作用了吧。

最后，《A Byte of Python》一書講裝飾器的時候提到了一篇博客：DRY Principles through Python Decorators 。有興趣的童鞋可以去閱讀以下。