9.5 可自定義屬性的裝飾器

問題

你想寫一個(gè)裝飾器來包裝一個(gè)函數(shù)，并且允許用戶提供參數(shù)在運(yùn)行時(shí)控制裝飾器行為。

解決方案

引入一個(gè)訪問函數(shù)，使用 nolocal 來修改內(nèi)部變量。然后這個(gè)訪問函數(shù)被作為一個(gè)屬性賦值給包裝函數(shù)。

from functools import wraps, partial
import logging
# Utility decorator to attach a function as an attribute of obj
def attach_wrapper(obj, func=None):
    if func is None:
        return partial(attach_wrapper, obj)
    setattr(obj, func.__name__, func)
    return func

def logged(level, name=None, message=None):
    '''
    Add logging to a function. level is the logging
    level, name is the logger name, and message is the
    log message. If name and message aren't specified,
    they default to the function's module and name.
    '''
    def decorate(func):
        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)

        # Attach setter functions
        @attach_wrapper(wrapper)
        def set_level(newlevel):
            nonlocal level
            level = newlevel

        @attach_wrapper(wrapper)
        def set_message(newmsg):
            nonlocal logmsg
            logmsg = newmsg

        return wrapper

    return decorate

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

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

下面是交互環(huán)境下的使用例子：

>>> import logging
>>> logging.basicConfig(level=logging.DEBUG)
>>> add(2, 3)
DEBUG:__main__:add
5
>>> # Change the log message
>>> add.set_message('Add called')
>>> add(2, 3)
DEBUG:__main__:Add called
5
>>> # Change the log level
>>> add.set_level(logging.WARNING)
>>> add(2, 3)
WARNING:__main__:Add called
5
>>>

討論

這一小節(jié)的關(guān)鍵點(diǎn)在于訪問函數(shù)(如 set_message() 和 set_level() )，它們被作為屬性賦給包裝器。每個(gè)訪問函數(shù)允許使用 nonlocal 來修改函數(shù)內(nèi)部的變量。

還有一個(gè)令人吃驚的地方是訪問函數(shù)會(huì)在多層裝飾器間傳播(如果你的裝飾器都使用了 @functools.wraps 注解)。例如，假設(shè)你引入另外一個(gè)裝飾器，比如9.2小節(jié)中的 @timethis ，像下面這樣：

@timethis
@logged(logging.DEBUG)
def countdown(n):
    while n > 0:
        n -= 1

你會(huì)發(fā)現(xiàn)訪問函數(shù)依舊有效：

>>> countdown(10000000)
DEBUG:__main__:countdown
countdown 0.8198461532592773
>>> countdown.set_level(logging.WARNING)
>>> countdown.set_message("Counting down to zero")
>>> countdown(10000000)
WARNING:__main__:Counting down to zero
countdown 0.8225970268249512
>>>

你還會(huì)發(fā)現(xiàn)即使裝飾器像下面這樣以相反的方向排放，效果也是一樣的：

@logged(logging.DEBUG)
@timethis
def countdown(n):
    while n > 0:
        n -= 1

還能通過使用lambda表達(dá)式代碼來讓訪問函數(shù)的返回不同的設(shè)定值：

@attach_wrapper(wrapper)
def get_level():
    return level

# Alternative
wrapper.get_level = lambda: level

一個(gè)比較難理解的地方就是對(duì)于訪問函數(shù)的首次使用。例如，你可能會(huì)考慮另外一個(gè)方法直接訪問函數(shù)的屬性，如下：

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

# Attach adjustable attributes
wrapper.level = level
wrapper.logmsg = logmsg
wrapper.log = log

這個(gè)方法也可能正常工作，但前提是它必須是最外層的裝飾器才行。如果它的上面還有另外的裝飾器(比如上面提到的 @timethis 例子)，那么它會(huì)隱藏底層屬性，使得修改它們沒有任何作用。而通過使用訪問函數(shù)就能避免這樣的局限性。

最后提一點(diǎn)，這一小節(jié)的方案也可以作為9.9小節(jié)中裝飾器類的另一種實(shí)現(xiàn)方法。