src/proxy_lite/logger.py

import logging
import sys
from typing import Literal
from uuid import uuid4

from rich.logging import RichHandler


class StructuredLogger(logging.Logger):
    def _log(
        self,
        level,
        msg,
        args,
        exc_info=None,
        extra=None,
        stack_info=False,
        stacklevel=1,
    ):
        if extra is None:
            extra = {}

        json_fields = {
            "logger_name": self.name,
            "message": msg % args if args else msg,
        }

        exc_type, exc_value, exc_traceback = sys.exc_info()
        if exc_type is not None:
            json_fields["exception_class"] = exc_type.__name__
            json_fields["exception_message"] = str(exc_value)

        json_fields.update(extra)
        super()._log(
            level,
            msg,
            args,
            exc_info,
            {"json_fields": json_fields},
            stack_info,
            stacklevel + 1,
        )


def create_logger(
    name: str,
    level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO",
    detailed_name: bool = False,
) -> logging.Logger:
    unique_name = f"{name}-{str(uuid4())[:8]}"
    logger = logging.getLogger(unique_name)
    logger.setLevel(level)
    handler = RichHandler(
        rich_tracebacks=True,
        markup=True,
        show_path=False,
        show_time=False,
        log_time_format="[%s]",
    )
    if detailed_name:
        handler.setFormatter(logging.Formatter("%(name)s:\n%(message)s\n------"))
    else:
        handler.setFormatter(logging.Formatter("%(message)s\n------"))
    logger.addHandler(handler)
    logger.propagate = False
    return logger


# Set StructuredLogger as the default logger class
logging.setLoggerClass(StructuredLogger)

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.propagate = True
handler = RichHandler(
    rich_tracebacks=True,
    markup=True,
    show_path=False,
    show_time=False,
)
logger.addHandler(handler)