Python Decorators: Streamlining Logging in Your Functions

Posted on by Kirin Choi

Python’s logging module is an excellent tool for tracking the activities of a script. For instance, if you want to log the start of a function, its parameters, return values, and completion, you can manually insert logging calls at appropriate places within the function. But what if you need to add these logs to every function you write? Copying and pasting similar logging calls into each function isn’t an elegant solution. As a software engineer, your instincts will likely tell you that there must be a more efficient and readable approach. This is where Python decorators come into play. Decorators provide a powerful and reusable way to enhance your functions with logging capabilities without cluttering your codebase.

Let’s say you have defined a function def greet(msg, name="World") as below:

def greet(msg, name="World"):
    return f"{msg}, {name}"

With logging module, you may add a log message for function parameters like this:

import logging

FORMAT = '[%(asctime)s] %(levelname)s - %(message)s'
logging.basicConfig(format=FORMAT, level=logging.INFO)
logger = logging.getLogger()

def greet(msg, name="World"):
    logger.info(f"Start greet({msg}, name={name})")
    return f"{msg}, {name}"

Let’s see if the logging module works correctly.

>>> greet("hello")
[2024-06-28 16:26:48,274] INFO - Start greet(hello, name=World)
'hello, World'
>>> greet("hi", name="Kirin")
[2024-06-28 16:27:18,948] INFO - Start greet(hi, name=Kirin)
'hi, Kirin'

While this approach works as expected, it becomes impractical when you have to add logging calls to every new function you define. This repetitive task is not only tedious but also undesirable. Here is the point where you need to consider adopting decorators, which can streamline the process and enhance the maintainability of your code. But what is a decorator?

Before I explain what a decorator is, first I’ll show how the code above will be changed to:

import logging

logger = logging.getLogger()

# The decorator will be defined here later.

@log
def greet(msg, name="World"):
    return f"{msg}, {name}"

The @log statement on line 7 is a decorator. By simply placing this decorator before the function, you can replace the manual logging calls with a more streamlined approach.

A decorator is a function which wraps another function and modifies its behavior. Before defining the decorator function, let’s explore what it means to wrap a function.

def log(func):
    def wrapper():
        print("Before the function call.")
        func()
        print("After the function call.")
    return wrapper

def tell_me_something():
    print("Something")

wrapped_function = log(tell_me_something)

Note that a function can be used as a parameter and a return value. You can also call the wrapped function like this:

>>> wrapped_function()
Before the function call.
Something
After the function call.

Instead of manually wrapping and calling a function, you can use a decorator by placing the wrapper function before the function definition with '@' symbol.

@log
def tell_me_others():
    print("Others")

>>> tell_me_others()
Before the function call.
Others
After the function call.

That’s it! Let’s apply this decorator to the greet() function and run it.

@log
def greet(msg, name="World"):
    return f"{msg}, {name}"

>>> greet("hi")
Traceback (most recent call last):
...
TypeError: wrapper() takes 0 positional arguments but 1 was given

It doesn’t work because our wrapper function isn’t passing positional and keyword arguments correctly. To fix this, let’s redefine the log decorator.

def log(func):
    def wrapper(*args, **kwargs):
        print("Before the function call.")
        func(*args, **kwargs)
        print("After the function call.")
    return wrapper

As you might know, *args and **kwargs are special symbols that allow you to pass an unspecified number of arguments to another function. With this fix, calling the greet() function will no longer raise an exception. But is this enough? No, there is still another problem. The return value of greet() is being swallowed by the @log decorator and isn’t passed to the caller. How can we resolve this? Simply passing the return value to the caller will do:

def log(func):
    def wrapper(*args, **kwargs):
        print("Before the function call.")
        result = func(*args, **kwargs)
        print("After the function call.")
        return result
    return wrapper

There is one last thing to consider. At this point, the Python interpreter will recognize the function by different name than greet(). To see this, try executing the following in the Python shell:

>>> greet
<function log.<locals>.wrapper at 0x1072a8ca0>

This is completely normal because it often doesn’t cause any issues. However, it doesn’t look neat. Fortunately, Python provides a good solution for this.

import functools

def log(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print("Before the function call.")
        result = func(*args, **kwargs)
        print("After the function call.")
        return result
    return wrapper

Redefine the @log decorator as shown above, redefine the function greet(), and then execute the following in the Python shell:

>>> greet
<function greet at 0x1072a8ca0>

Now it looks better. From now on, all you need to do is replace the print function with the logging module and tailor log messages as you wish. Here is my version of the @log decorator:

import logging
import functools

FORMAT = '[%(asctime)s] %(levelname)s - %(message)s'
logging.basicConfig(format=FORMAT, level=logging.INFO)
logger = logging.getLogger(__name__)

def log(func):
    def to_str(v):
        if isinstance(v, pd.DataFrame):
            return f'DataFrame{v.shape}'
        elif isinstance(v, list):
            return f'list({len(v)})'
        elif isinstance(v, pymssql.Connection):
            return 'pymssql.Connection'
        elif isinstance(v, datetime.datetime):
            return f"{v.strftime('%Y-%m-%d %H:%M:%S')}"
        else:
            return repr(v)

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        args_repr = [to_str(a) for a in args]
        kwargs_repr = [f"{k}={to_str(v)}" for k, v in kwargs.items()]
        signature = ", ".join(args_repr + kwargs_repr)
        logger.info(f"Start {func.__name__}({signature})")
        result = func(*args, **kwargs)
        logger.info(f"End {func.__name__}() => {to_str(result)}")
        return result

    return wrapper

The function to_str(v) is defined to convert values to a proper string format. In most cases, repr() or str() will suffice. However, sometimes these functions produce messy outputs that are difficult to read. For instance, using repr() on a pandas.DataFrame1) variable results in a huge matrix form. For my scripts, pandas.DataFrame.shape is usually sufficient.

Although it looks more complex, the wrapper() function is almost identical to the original. The only difference is that my version prints out the function’s name, parameter values, and return values.

In conclusion, using decorators for logging in Python can greatly enhance the readability and maintainability of your code. By encapsulating the logging logic within a decorator, you can easily add comprehensive logging to any function without cluttering your codebase. I hope this article has provided you with a clear understanding of how to implement and utilize decorators for logging in your own projects. Happy coding!

References

1) pandas is a well-known open source data analysis and manipulation tool for Python. The DataFrame class, defined in the pandas module, is specifically designed and optimized for handling large tabular data.

Leave a comment