Python decorator for logging the input and output of a class method

log_io.py

import inspect
import logging
from functools import wraps
from typing import Callable


def log_io(log_input=False, log_output=False) -> Callable:
    """
    Logs the input/output of the decorated method.
    Must be explicitly called with input=True and/or output=True
    >>> @log_io(log_input=True, log_output=True)
    >>> @log_io(log_input=True)
    >>> @log_io(log_output=True)
    >>>
    >>> class Foo:
    >>>     @log_io(log_input=True, log_output=True)
    >>>     def bar(self, a, b):
    >>>         return a + b
    """
    def decorator(func):
        # get the caller method, for being able to preserve 
        # the signature of the decorated method
        caller = inspect.getframeinfo(inspect.stack()[1][0])

        @wraps(func)
        def inner(*args, **kwargs):
            if not any([log_input, log_output]):
                return func(*args, **kwargs)
            instance = args[0]  # class instance aka self
            # preserve the original function name, module and file_no
            # to not point to the decorator location
            extra = {
                "func_name": func.__name__,
                "line_no": caller.lineno,
                "path": caller.filename,
            }
            logger = logging.getLogger(func.__module__)

            if log_input:
                logger.info(f"{instance!r} INPUT: {args[1:]}, {kwargs}", extra=extra)

            if log_output:
                try:
                    output = func(*args, **kwargs)
                    logger.info(f"{instance!r} OUTPUT: {output}", extra=extra)
                    return output
                except Exception as e:
                    logger.exception(f"{instance!r} OUTPUT: {e}", extra=extra)
                    raise e

            return func(*args, **kwargs)
        return inner
    return decorator