使用裝飾器將參數添加到裝飾函數

Question

假設我有一個函數foo在執行期間記錄一些信息：

import logging

logging.basicConfig()

def foo(x):
    for i in range(10):
        logging.info(f"This is iteration number {i}.")
        x += i
    return x

是否可以創建一個裝飾器log_or_not以便我可以調用我的裝飾函數foo並添加如下參數：

foo(level=logging.INFO)

現在我只有這個，一個將日志記錄級別設置為 INFO 的裝飾器，執行該函數然后將日志記錄級別返回到原始級別：

def log_or_not(func):
    def wrap_function(*args, **kwargs):
        original_level = logging.root.level
        print(original_level)
        logging.getLogger().setLevel(logging.INFO)
        y = func(*args, **kwargs)
        logging.getLogger().setLevel(original_level)
        return y
    return wrap_function

@log_or_not
def foo(x):
    for i in range(10):
        logging.info(f"This is iteration number {i}.")
        x += i
    return x

Answer 1

將額外參數添加到您的wrap_function ：

import logging

def log_or_not(func):
    def wrap_function(*args, level=logging.INFO, **kwargs):
        original_level = logging.root.level
        print(original_level)
        logging.getLogger().setLevel(level)
        try:
            return func(*args, **kwargs)
        finally:
            logging.getLogger().setLevel(original_level)
    return wrap_function

使用try/finally可以保證您在調用func后重置原始日志記錄級別，即使它引發了異常。

由於wrap_function取代了裝飾函數，您現在可以執行以下操作：

@log_or_not
def foo(x):
    for i in range(10):
        logging.info(f"This is iteration number {i}.")
        x += i
    return x

foo(0, level=logging.ERROR)  # logs with level ERROR
foo(0)                       # logs with level INFO (default)

Answer 2

我會使用上下文管理器來動態修改日志記錄級別，而讓函數本身保持獨立。

import contextlib

# No try needed; the context manager decorator already 
# ensures that the post-yield code is executed before  
# propagating any exceptions.
@contextlib.contextmanager
def change_level(lvl):
    original = logging.root.level
    logging.getLogger().setLevel(lvl)
    yield
    logging.getLogger().setLevel(original)


# instead of foo(0, level=logging.ERROR)
with change_level(logging.INFO):
    foo(0)