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

問(wèn)題

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

解決方案

可以使用關(guān)鍵字參數(shù)來(lái)給被包裝函數(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
>>>

討論

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

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')

那么你可以將其重構(gòu)成這樣：

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

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

還有一個(gè)難點(diǎn)就是如何去處理被添加的參數(shù)與被包裝函數(shù)參數(shù)直接的名字沖突。例如，如果裝飾器 @optional_debug 作用在一個(gè)已經(jīng)擁有一個(gè) debug 參數(shù)的函數(shù)上時(shí)會(huì)有問(wèn)題。這里我們?cè)黾恿艘徊矫謾z查。

上面的方案還可以更完美一點(diǎn)，因?yàn)榫鞯某绦騿T應(yīng)該發(fā)現(xiàn)了被包裝函數(shù)的函數(shù)簽名其實(shí)是錯(cuò)誤的。例如：

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

通過(guò)如下的修改，可以解決這個(gè)問(wèn)題：

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

通過(guò)這樣的修改，包裝后的函數(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é)獲取更多關(guān)于函數(shù)簽名的信息。