使用装饰器将参数添加到装饰函数

Question

Let's say I have a function foo that logs some info during execution:

import logging

logging.basicConfig()

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

Is it possible to create a decorator log_or_not so that I can call my decorated function foo with an added argument like this:

foo(level=logging.INFO)

For now all I have is this, a decorator that sets the logging level to INFO, executes the function then returns the logging level to the original level:

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

Add the extra parameter to your 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

Using try/finally guarantees that you'll reset the original logging level after calling func even if it raises an exception.

Since wrap_function takes the place of the decorated function, you can now do:

@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

I would use a context manager to modify logging levels on the fly, and leave the function itself alone.

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)