9.11 裝飾器為被包裝函數(shù)增加參數(shù)

問題

你想在裝飾器中給被包裝函數(shù)增加額外的參數(shù)，但是不能影響這個函數(shù)現(xiàn)有的調(diào)用規(guī)則。

解決方案

可以使用關鍵字參數(shù)來給被包裝函數(shù)增加額外參數(shù)?？紤]下面的裝飾器：

from functools import wraps

def optional_debug(func):
    @wraps(func)
    def wrapper(*args, debug=False, **kwargs):
        if debug:
            print('Calling', func.__name__)
        return func(*args, **kwargs)

    return wrapper

>>> @optional_debug
... def spam(a,b,c):
... print(a,b,c)
...
>>> spam(1,2,3)
1 2 3
>>> spam(1,2,3, debug=True)
Calling spam
1 2 3
>>>

討論

通過裝飾器來給被包裝函數(shù)增加參數(shù)的做法并不常見。盡管如此，有時候它可以避免一些重復代碼。例如，如果你有下面這樣的代碼：

def a(x, debug=False):
    if debug:
        print('Calling a')

def b(x, y, z, debug=False):
    if debug:
        print('Calling b')

def c(x, y, debug=False):
    if debug:
        print('Calling c')

那么你可以將其重構成這樣：

from functools import wraps
import inspect

def optional_debug(func):
    if 'debug' in inspect.getargspec(func).args:
        raise TypeError('debug argument already defined')

    @wraps(func)
    def wrapper(*args, debug=False, **kwargs):
        if debug:
            print('Calling', func.__name__)
        return func(*args, **kwargs)
    return wrapper

@optional_debug
def a(x):
    pass

@optional_debug
def b(x, y, z):
    pass

@optional_debug
def c(x, y):
    pass

這種實現(xiàn)方案之所以行得通，在于強制關鍵字參數(shù)很容易被添加到接受 *args 和 **kwargs 參數(shù)的函數(shù)中。通過使用強制關鍵字參數(shù)，它被作為一個特殊情況被挑選出來，并且接下來僅僅使用剩余的位置和關鍵字參數(shù)去調(diào)用這個函數(shù)時，這個特殊參數(shù)會被排除在外。也就是說，它并不會被納入到 **kwargs 中去。

還有一個難點就是如何去處理被添加的參數(shù)與被包裝函數(shù)參數(shù)直接的名字沖突。例如，如果裝飾器 @optional_debug 作用在一個已經(jīng)擁有一個 debug 參數(shù)的函數(shù)上時會有問題。這里我們增加了一步名字檢查。

上面的方案還可以更完美一點，因為精明的程序員應該發(fā)現(xiàn)了被包裝函數(shù)的函數(shù)簽名其實是錯誤的。例如：

>>> @optional_debug
... def add(x,y):
...     return x+y
...
>>> import inspect
>>> print(inspect.signature(add))
(x, y)
>>>

通過如下的修改，可以解決這個問題：

from functools import wraps
import inspect

def optional_debug(func):
    if 'debug' in inspect.getargspec(func).args:
        raise TypeError('debug argument already defined')

    @wraps(func)
    def wrapper(*args, debug=False, **kwargs):
        if debug:
            print('Calling', func.__name__)
        return func(*args, **kwargs)

    sig = inspect.signature(func)
    parms = list(sig.parameters.values())
    parms.append(inspect.Parameter('debug',
                inspect.Parameter.KEYWORD_ONLY,
                default=False))
    wrapper.__signature__ = sig.replace(parameters=parms)
    return wrapper

通過這樣的修改，包裝后的函數(shù)簽名就能正確的顯示 debug 參數(shù)的存在了。例如：

>>> @optional_debug
... def add(x,y):
...     return x+y
...
>>> print(inspect.signature(add))
(x, y, *, debug=False)
>>> add(2,3)
5
>>>

參考9.16小節(jié)獲取更多關于函數(shù)簽名的信息。