2
\$\begingroup\$

I want to write a decorator in Python 3.11 that adds some basic logging to functions.

I want to use the decorator without any keyword arguments:

@add_logging

In that case it should log with a default logging level of logging.DEBUG.

I also want to use the decorator with a keyword argument that then specifies a logging level:

@add_logging(logging_level=logging.ERROR)

And all of that should be typed to its fullest extent.

I figured out how to do that using the @overload decorator:

P = ParamSpec("P")
R = TypeVar("R")


@overload
def add_logging(function: Callable[P, R]) -> Callable[P, R]:
    ...


@overload
def add_logging(*, logging_level: int = DEBUG) -> Callable[[Callable[P, R]], Callable[P, R]]:
    ...


def add_logging(
    function: Optional[Callable[P, R]] = None, *, logging_level: int = DEBUG
) -> Callable[[Callable[P, R]], Callable[P, R]] | Callable[P, R]:
    """A type-safe decorator to add logging to a function.

    """

    def wrapper(wrapped_function: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
        module_name: str = wrapped_function.__module__
        logger: Logger = getLogger(module_name)

        if not logger.isEnabledFor(logging_level):
            return wrapped_function(*args, **kwargs)

        filename: str = wrapped_function.__code__.co_filename
        first_line_no: int = wrapped_function.__code__.co_firstlineno

        logger.handle(
            LogRecord(
                name=module_name,
                level=logging_level,
                pathname=filename,
                lineno=first_line_no,
                msg="Call *%r **%r",
                args=(args, kwargs),
                exc_info=None,
                func=wrapped_function.__qualname__,
            )
        )

        try:
            result = wrapped_function(*args, **kwargs)
        except Exception as exception:
            logger.handle(
                LogRecord(
                    name=module_name,
                    level=logging_level,
                    pathname=filename,
                    lineno=first_line_no,
                    msg="",
                    args=None,
                    exc_info=(type(exception), exception, None),
                    func=wrapped_function.__qualname__,
                )
            )
            raise

        logger.handle(
            LogRecord(
                name=module_name,
                level=logging_level,
                pathname=filename,
                lineno=first_line_no,
                msg="Return %r",
                args=(result,),
                exc_info=None,
                func=wrapped_function.__qualname__,
            )
        )
        return result

    # Without arguments, `function` is passed directly to the decorator
    if function is not None:
        if not callable(function):
            raise TypeError(f"Expected positional parameter of type callable, but found type {type(function)} instead.")
        return wraps(function)(partial(wrapper, function))

    # With arguments, we need to return a function that accepts the function
    def decorator(function_with_args: Callable[P, R]) -> Callable[P, R]:
        return wraps(function_with_args)(partial(wrapper, function_with_args))

    return decorator

This is fully accepted by mypy and works as described above, that is, with and without the keyword argument logging_level. Here's an example output of a function decorated with this:

[2023-05-01 00:08:19 +0100] [__main__] [add_two_ints] [DEBUG] Call *(5, 7) **{}
[2023-05-01 00:08:19 +0100] [__main__] [add_two_ints] [DEBUG] Return 12

The reason I used LogRecord instead of just logging.debug/info/etc is the flexibility in terms of the logging level and the ability to explicitly set the name of the function that is wrapped as this is part of my logging strategy.

However, the code looks quite bloated. Am I missing something here? Is there an easier, more elegant, more obvious way to achieve this?

Please note that I am looking for fully typed code examples, as this was one of the challenges I faced when trying to solve this.

\$\endgroup\$
3
  • \$\begingroup\$ You wanted two implementations of add_logging, why do you have three? \$\endgroup\$ May 2, 2023 at 6:00
  • \$\begingroup\$ Welcome to Code Review! I changed the title so that it describes what the code does per site goals: "State what your code does in your title, not your main concerns about it.". Please check that I haven't misrepresented your code, and correct it if I have. \$\endgroup\$ May 2, 2023 at 6:48
  • \$\begingroup\$ @sbottingota Well, the way I look at it there is only one implementation and that is overloaded. But please correct me if I am wrong. \$\endgroup\$
    – uli
    May 3, 2023 at 7:10

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.