python - 使用装饰器为装饰函数添加一个参数

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(x=x, 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)