src.common.base_log 源代码

import inspect
import json
from loguru import logger
import os
import sys
from typing import Optional, Union

LOG_COLORS = {
    "DEBUG": "\033[1;36m",  # CYAN
    "INFO": "\033[1;32m",  # GREEN
    "WARNING": "\033[1;33m",  # YELLOW
    "ERROR": "\033[1;31m",  # RED
    "CRITICAL": "\033[1;31m",  # RED
    "EXCEPTION": "\033[1;31m",  # RED
}
COLOR_RESET = "\033[1;0m"



class Log:
    # 静态标记：确保只移除一次默认 handler
    _default_handler_removed = False

    def __init__(
        self,
        filename: Optional[str] = None,
        cmdlevel: str = "DEBUG",
        filelevel: str = "INFO",
        backup_count: int = 7,  # 默认保留7天/7个文件
        limit: Union[int, str] = "20 MB",  # 支持字符串格式
        when: Optional[str] = None,
        colorful: bool = True,
        compression: Optional[str] = None,  # 新增压缩功能
        is_backtrace: bool = True,
    ):
        # 只移除 loguru 默认的 stderr handler (ID 0)，避免删除其他模块已添加的 handlers
        if not Log._default_handler_removed:
            try:
                logger.remove(0)  # 只移除默认 handler
            except ValueError:
                pass  # 默认 handler 可能已被移除
            Log._default_handler_removed = True

        self.is_backtrace = is_backtrace
        self.logger = logger.bind(task=filename)
        # 设置日志文件路径
        if filename is None:
            filename = getattr(sys.modules["__main__"], "__file__", "log.py")
            filename = os.path.basename(filename.replace(".py", ".log"))

        # 确保日志目录存在
        log_dir = os.path.abspath(os.path.dirname(filename))
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)

        # 控制台输出配置
        self.logger.add(
            sys.stderr,
            level=cmdlevel,
            format=self._formatter,
            colorize=colorful,
            backtrace=True,
            enqueue=True,
            filter=lambda record: record["extra"].get("task") == filename,
        )

        # 文件输出配置
        rotation_config = self._get_rotation_config(when, limit)
        self.logger.add(
            filename,
            level=filelevel,
            format=self._formatter,
            backtrace=True,
            rotation=rotation_config,
            retention=f"{backup_count} days",
            compression=compression,
            enqueue=True,
            filter=lambda record: record["extra"].get("task") == filename,
        )

    def _formatter(self, record):
        # 处理消息内容
        message = str(record["message"]).replace("{", "{{").replace("}", "}}")
        if isinstance(message, (dict, list)):
            try:
                message = json.dumps(message, ensure_ascii=False)
            except TypeError:
                message = str(message)

        # 处理调用栈和颜色
        if self.is_backtrace:
            frame = inspect.currentframe()
            while frame:
                if (
                    "loguru" not in frame.f_code.co_filename
                    and "base_log.py" not in frame.f_code.co_filename
                ):
                    break
                frame = frame.f_back
            file_info = (
                f"[{os.path.basename(frame.f_code.co_filename)}:{frame.f_code.co_name}:{frame.f_lineno}]"
                if frame
                else f"[{record['file']}:{record['function']}:{record['line']}]"
            )
        else:
            file_info = f"[{record['file']}:{record['line']}]"

        level_color = LOG_COLORS.get(record["level"].name, "")
        return (
            f"{level_color}[{record['time'].strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}] "
            + file_info
            + f"[{record['level']}] {message}{COLOR_RESET}\n"
        )

    def _get_rotation_config(self, when: Optional[str], limit: Union[int, str]):
        if when:  # 时间轮转
            return when  # "D"（天）、"H"（小时）、"midnight"等
        else:  # 大小轮转
            if isinstance(limit, int):
                return f"{limit / 1024 / 1024} MB"
            return limit  # 直接支持"10 MB"、"1 GB"等字符串格式



    @staticmethod
    def set_logger(**kwargs) -> bool:
        """For backward compatibility."""
        return True




    def debug(self, *args, **kwargs):
        self.logger.debug(*args, **kwargs)




    def info(self, *args, **kwargs):
        self.logger.info(*args, **kwargs)




    def warning(self, *args, **kwargs):
        self.logger.warning(*args, **kwargs)




    def error(self, *args, **kwargs):
        self.logger.error(*args, **kwargs)




    def critical(self, *args, **kwargs):
        self.logger.critical(*args, **kwargs)




    def set_config(
        self,
        filename: Optional[str] = None,
        cmdlevel: str = "DEBUG",
        filelevel: str = "INFO",
        backup_count: int = 7,
        limit: Union[int, str] = "20 MB",
        when: Optional[str] = None,
        colorful: bool = True,
        compression: Optional[str] = None,
    ):
        """动态修改日志配置"""
        # 移除旧的handler
        self.logger.remove()
        
        # 重新添加handler
        if filename is None:
            filename = getattr(sys.modules["__main__"], "__file__", "log.py")
            filename = os.path.basename(filename.replace(".py", ".log"))

        # 确保日志目录存在
        log_dir = os.path.abspath(os.path.dirname(filename))
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)

        # 控制台输出配置
        self.logger.add(
            sys.stderr,
            level=cmdlevel,
            format=self._formatter,
            colorize=colorful,
            backtrace=True,
            enqueue=True,
            filter=lambda record: record["extra"].get("task") == filename,
        )

        # 文件输出配置
        rotation_config = self._get_rotation_config(when, limit)
        self.logger.add(
            filename,
            level=filelevel,
            format=self._formatter,
            backtrace=True,
            rotation=rotation_config,
            retention=f"{backup_count} days",
            compression=compression,
            enqueue=True,
            filter=lambda record: record["extra"].get("task") == filename,
        )




    def exception(self, *args, **kwargs):
        self.logger.exception(*args, **kwargs)