Use decorator to add an argument to the decorated function

July 19, 2022

Let’s say I have a function:

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

>Solution :

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)