json_logger.py

"""Logging for the project."""

import json
import logging


class JsonFormatter(logging.Formatter):
    """Formatter that outputs JSON strings after parsing the LogRecord.

    Notes:
        From here: https://stackoverflow.com/a/70223539/21188807
    """

    def __init__(
        self,
        fmt_dict: dict = None,
        time_format: str = "%Y-%m-%dT%H:%M:%S",
        msec_format: str = "%s.%03dZ",
    ):
        """Initialize the JSONFormatter class."""
        self.fmt_dict = fmt_dict if fmt_dict is not None else {"message": "message"}
        self.default_time_format = time_format
        self.default_msec_format = msec_format
        self.datefmt = None

    def uses_time(self) -> bool:
        """Look for the attribute in the format dict values instead of the fmt string."""
        return "asctime" in self.fmt_dict.values()

    def format_message(self, record) -> dict:
        """Return a dictionary of the relevant LogRecord attributes instead of a string.

        Raises:
            KeyError is raised if an unknown attribute is provided in the fmt_dict.

        """
        return {
            fmt_key: record.__dict__[fmt_val]
            for fmt_key, fmt_val in self.fmt_dict.items()
        }

    def format(self, record) -> str:
        """Format the record as json for logging."""
        record.message = record.getMessage()

        if self.uses_time():
            record.asctime = self.formatTime(record, self.datefmt)

        message_dict = self.format_message(record)

        if record.exc_info:
            # Cache the traceback text to avoid converting it multiple times
            # (it's constant anyway)
            if not record.exc_text:
                record.exc_text = self.formatException(record.exc_info)

        if record.exc_text:
            message_dict["exc_info"] = record.exc_text

        if record.stack_info:
            message_dict["stack_info"] = self.formatStack(record.stack_info)

        return json.dumps(message_dict, default=str)


def get_logger(name) -> logging.Logger:
    """Get a default logger with name `name`."""
    logger = logging.getLogger(name)

    if not logger.hasHandlers():
        handler = logging.StreamHandler()
        formatter = JsonFormatter(
            {
                "time": "asctime",
                "severity": "levelname",
                "message": "message",
                "log": "name",
                "filename": "filename",
                "line": "lineno",
                "pathname": "pathname",
            }
        )

        handler.setFormatter(formatter)

        logger.addHandler(handler)
        logger.setLevel(logging.INFO)

    return logger