graphgen/utils/log.py

import contextvars
import logging
import os
from logging.handlers import RotatingFileHandler
from typing import Any

from rich.logging import RichHandler


def set_logger(
    log_file: str,
    name: str,
    file_level: int = logging.DEBUG,
    console_level: int = logging.INFO,
    *,
    if_stream: bool = True,
    max_bytes: int = 50 * 1024 * 1024,  # 50 MB
    backup_count: int = 5,
    force: bool = False,
):

    current_logger = logging.getLogger(name)
    if current_logger.hasHandlers() and not force:
        return current_logger

    if force:
        current_logger.handlers.clear()

    current_logger.setLevel(
        min(file_level, console_level)
    )  # Set to the lowest level to capture all logs
    current_logger.propagate = False

    if log_file:
        os.makedirs(os.path.dirname(log_file), exist_ok=True)

    if if_stream:
        console = RichHandler(
            level=console_level, show_path=False, rich_tracebacks=True
        )
        console.setFormatter(logging.Formatter("%(message)s"))
        current_logger.addHandler(console)

    file_handler = RotatingFileHandler(
        log_file,
        maxBytes=max_bytes,
        backupCount=backup_count,
        encoding="utf-8",
    )
    file_handler.setLevel(file_level)
    file_handler.setFormatter(
        logging.Formatter(
            "[%(asctime)s] %(levelname)s [%(name)s:%(filename)s:%(lineno)d] %(message)s",
            datefmt="%y-%m-%d %H:%M:%S",
        )
    )
    current_logger.addHandler(file_handler)
    return current_logger


CURRENT_LOGGER_VAR = contextvars.ContextVar("current_logger")


def get_current_logger() -> logging.Logger:
    current_logger = CURRENT_LOGGER_VAR.get()
    if not current_logger:
        raise RuntimeError("No logger is set in the current context.")
    return current_logger


class ContextAwareLogger:
    @staticmethod
    def _get_logger() -> logging.Logger:
        return get_current_logger()

    def debug(self, msg: object, *args: Any, **kwargs: Any) -> None:
        self._get_logger().debug(msg, *args, **kwargs)

    def info(self, msg: object, *args: Any, **kwargs: Any) -> None:
        self._get_logger().info(msg, *args, **kwargs)

    def warning(self, msg: object, *args: Any, **kwargs: Any) -> None:
        self._get_logger().warning(msg, *args, **kwargs)

    def error(self, msg: object, *args: Any, **kwargs: Any) -> None:
        self._get_logger().error(msg, *args, **kwargs)

    def exception(self, msg: object, *args: Any, **kwargs: Any) -> None:
        self._get_logger().exception(msg, *args, **kwargs)

    def critical(self, msg: object, *args: Any, **kwargs: Any) -> None:
        self._get_logger().critical(msg, *args, **kwargs)

    def log(self, level: int, msg: object, *args: Any, **kwargs: Any) -> None:
        self._get_logger().log(level, msg, *args, **kwargs)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._get_logger(), name)


logger = ContextAwareLogger()