9.6 帶可選參數(shù)的裝飾器

問題

你想寫一個(gè)裝飾器，既可以不傳參數(shù)給它，比如 @decorator ，也可以傳遞可選參數(shù)給它，比如 @decorator(x,y,z) 。

解決方案

下面是9.5小節(jié)中日志裝飾器的一個(gè)修改版本：

from functools import wraps, partial
import logging

def logged(func=None, *, level=logging.DEBUG, name=None, message=None):
    if func is None:
        return partial(logged, level=level, name=name, message=message)

    logname = name if name else func.__module__
    log = logging.getLogger(logname)
    logmsg = message if message else func.__name__

    @wraps(func)
    def wrapper(*args, **kwargs):
        log.log(level, logmsg)
        return func(*args, **kwargs)

    return wrapper

# Example use
@logged
def add(x, y):
    return x + y

@logged(level=logging.CRITICAL, name='example')
def spam():
    print('Spam!')

可以看到，@logged 裝飾器可以同時(shí)不帶參數(shù)或帶參數(shù)。

討論

這里提到的這個(gè)問題就是通常所說的編程一致性問題。當(dāng)我們使用裝飾器的時(shí)候，大部分程序員習(xí)慣了要么不給它們傳遞任何參數(shù)，要么給它們傳遞確切參數(shù)。其實(shí)從技術(shù)上來講，我們可以定義一個(gè)所有參數(shù)都是可選的裝飾器，就像下面這樣：

@logged()
def add(x, y):
    return x+y

但是，這種寫法并不符合我們的習(xí)慣，有時(shí)候程序員忘記加上后面的括號(hào)會(huì)導(dǎo)致錯(cuò)誤。這里我們向你展示了如何以一致的編程風(fēng)格來同時(shí)滿足沒有括號(hào)和有括號(hào)兩種情況。

為了理解代碼是如何工作的，你需要非常熟悉裝飾器是如何作用到函數(shù)上以及它們的調(diào)用規(guī)則。對(duì)于一個(gè)像下面這樣的簡單裝飾器：

# Example use
@logged
def add(x, y):
    return x + y

這個(gè)調(diào)用序列跟下面等價(jià)：

def add(x, y):
    return x + y

add = logged(add)

這時(shí)候，被裝飾函數(shù)會(huì)被當(dāng)做第一個(gè)參數(shù)直接傳遞給 logged 裝飾器。因此，logged() 中的第一個(gè)參數(shù)就是被包裝函數(shù)本身。所有其他參數(shù)都必須有默認(rèn)值。

而對(duì)于一個(gè)下面這樣有參數(shù)的裝飾器：

@logged(level=logging.CRITICAL, name='example')
def spam():
    print('Spam!')

調(diào)用序列跟下面等價(jià)：

def spam():
    print('Spam!')
spam = logged(level=logging.CRITICAL, name='example')(spam)

初始調(diào)用 logged() 函數(shù)時(shí)，被包裝函數(shù)并沒有傳遞進(jìn)來。因此在裝飾器內(nèi)，它必須是可選的。這個(gè)反過來會(huì)迫使其他參數(shù)必須使用關(guān)鍵字來指定。并且，但這些參數(shù)被傳遞進(jìn)來后，裝飾器要返回一個(gè)接受一個(gè)函數(shù)參數(shù)并包裝它的函數(shù)(參考9.5小節(jié))。為了這樣做，我們使用了一個(gè)技巧，就是利用 functools.partial 。它會(huì)返回一個(gè)未完全初始化的自身，除了被包裝函數(shù)外其他參數(shù)都已經(jīng)確定下來了?？梢詤⒖?.8小節(jié)獲取更多 partial() 方法的知識(shí)。