問題

你想寫一個裝飾器來包裝一個函數，并且允許用戶提供參數在運行時控制裝飾器行為。

解決方案

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

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é)的關鍵點在于訪問函數(如 set_message() 和 set_level() )，它們被作為屬性賦給包裝器。每個訪問函數允許使用 nonlocal 來修改函數內部的變量。

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

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

你會發(fā)現訪問函數依舊有效：

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

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

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

還能通過使用lambda表達式代碼來讓訪問函數的返回不同的設定值：

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

# Alternative
wrapper.get_level = lambda: level

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

@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

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

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