logger decorator for async app

log_extra.py

"""
Adds async context to a Python logger.
"""
import asyncio
import functools
import inspect
import logging
from contextvars import ContextVar
from typing import Any, Callable

_EXTRA_TYPE = dict[str, Any]
_EXTRA_VAR = ContextVar("extra", default=dict())


class Logger:
    """
    Wrapper for a python :py:class:`logging.Logger`.

    This wrapper reads the current local context and emits
    them for each log message.
    """

    _original_get = None

    def __init__(self, name: str | None = None):
        self.base_logger = self.__class__._original_get(name)

    @classmethod
    def setup(cls):
        cls._original_get = logging.getLogger
        logging.getLogger = Logger

    def _msg(self, func: Callable, msg, *args, **kwargs):
        """Log with our extra values,"""
        extra = _EXTRA_VAR.get() | kwargs.pop("extra", {})
        # Because we wrap a logging.Logger instance through 2 layers
        # of redirection, we need to add 2 to the logger's stacklevel
        # so we correctly log the logging statement's line number and function name
        original_stacklevel = kwargs.pop("stacklevel", 1)
        stacklevel = original_stacklevel + 2
        return func(msg, *args, extra=extra, stacklevel=stacklevel, **kwargs)

    def __getattr__(self, item):
        attr = getattr(self.base_logger, item)
        if attr in ("debug", "info", "warning", "error", "exception", "critical"):

            def _log(*args, **kwargs):
                self._msg(attr, *args, **kwargs)

            return _log
        return attr


class log_extra:
    def __init__(self, extra=None, *_, **kwargs):
        self.old_extra = dict()
        self.new_extra = (extra or {}) | kwargs

    def __enter__(self):
        self.old_extra = _EXTRA_VAR.get()
        _EXTRA_VAR.set(self.old_extra | self.new_extra)

    def __exit__(self, *args, **kwds):
        _EXTRA_VAR.set(self.old_extra)

    def __call__(self, f: Callable | type):
        if inspect.isclass(f):
            for attr in f.__dict__:  # there's propably a better way to do this
                if callable(getattr(f, attr)):
                    setattr(f, attr, self(getattr(f, attr)))
            return f
        elif inspect.iscoroutinefunction(f):

            @functools.wraps(f)
            async def actual_decorator(*args, **kwargs):
                with self:
                    return await f(*args, **kwargs)

        else:

            @functools.wraps(f)
            def actual_decorator(*args, **kwargs):
                with self:
                    return f(*args, **kwargs)

        return actual_decorator

tests.py

import asyncio
import logging

from .log_extra import log_extra, _EXTRA_VAR


async def handler():
    print(_EXTRA_VAR.get())


async def handler2():
    with log_extra(id="2"):
        await handler()


@log_extra(id="3")
async def handler3():
    await handler()


@log_extra(class_name="X")
class X:
    def x(self):
        print(_EXTRA_VAR.get())

    async def y(self):
        print(_EXTRA_VAR.get())


async def task():
    Logger(__name__).error("My message")


@log_extra(id=1)
async def main():
    await handler()
    await handler2()
    await handler3()
    await handler()
    X().x()
    await X().y()
    _task = asyncio.create_task(task())
    while not _task.done():
        await asyncio.sleep(1)


Logger.setup()
logger = logging.getLogger(__name__)
logger.error("My message 2")

asyncio.run(main())