logger.py

import logging
from datetime import datetime
from typing import TYPE_CHECKING, Any, MutableMapping

from pythonjsonlogger import jsonlogger


class CustomJsonFormatter(jsonlogger.JsonFormatter):
    def add_fields(
        self,
        log_record: dict[str, Any],
        record: logging.LogRecord,
        message_dict: dict[str, Any],
    ) -> None:
        super(CustomJsonFormatter, self).add_fields(
            log_record, record, message_dict
        )
        if not log_record.get("timestamp"):
            # this doesn't use record.created, so it is slightly off
            now = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ")
            log_record["timestamp"] = now
        if log_record.get("level"):
            log_record["level"] = log_record["level"].upper()
        else:
            log_record["level"] = record.levelname


formatter = CustomJsonFormatter()


# Create a custom logger
logger = logging.getLogger(name="tuhag")

# Create handlers
c_handler = logging.StreamHandler()
f_handler = logging.FileHandler("file.log")
c_handler.setLevel(logging.WARNING)
f_handler.setLevel(logging.DEBUG)

# Create formatters and add it to handlers
c_handler.setFormatter(formatter)
f_handler.setFormatter(formatter)

# Add handlers to the logger
logger.addHandler(c_handler)
logger.addHandler(f_handler)


if TYPE_CHECKING:
    LoggerAdapterBase = logging.LoggerAdapter[logging.Logger]
else:
    LoggerAdapterBase = logging.LoggerAdapter


class CustomLoggerAdapter(LoggerAdapterBase):
    def process(
        self, msg: str, kwargs: MutableMapping[str, Any]
    ) -> tuple[str, MutableMapping[str, Any]]:
        """
        The `extra` in `LoggerAdapter` takes precedence over the `extra` argument in the logging call's argument.
        This CustomLoggerAdapter merges the 2 dicts together
        """
        if "extra" in kwargs:
            # this merges the two extras together
            # self.extra can be None
            merged = {**(self.extra or {}), **kwargs["extra"]}
            kwargs["extra"] = merged
        else:
            # when the logging call does not have the `extra` argument
            kwargs["extra"] = self.extra
        return msg, kwargs