Spaces:
Running
Running
File size: 4,119 Bytes
583741e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
import logging
from enum import IntEnum
from typing import Any, Optional
from rich.console import Console, Group
from rich.panel import Panel
from rich.rule import Rule
from rich.syntax import Syntax
from rich.table import Table
from rich.tree import Tree
from rich.logging import RichHandler
from src.utils import Singleton
YELLOW_HEX = "#d4b702"
class LogLevel(IntEnum):
CRITICAL = logging.CRITICAL
FATAL = logging.FATAL
ERROR = logging.ERROR
WARNING = logging.WARNING
WARN = logging.WARN
INFO = logging.INFO
DEBUG = logging.DEBUG
class Logger(logging.Logger, metaclass=Singleton):
def __init__(self, name="logger", level=logging.INFO):
# Initialize the parent class
super().__init__(name, level)
# Define a formatter for log messages
self.formatter = logging.Formatter(
fmt="%(asctime)s - %(name)s:%(levelname)s - %(filename)s:%(lineno)s - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
def init_logger(self, config, level: int = LogLevel.INFO):
"""
Initialize the logger with a file path and optional main process check.
Args:
log_path (str): The log file path.
level (int, optional): The logging level. Defaults to logging.INFO.
accelerator (Accelerator, optional): Accelerator instance to determine the main process.
"""
log_path = config.log_path
self.handlers.clear()
self.console = Console(
width=None,
markup=True,
color_system="truecolor",
force_terminal=True
)
rich_handler = RichHandler(
console=self.console,
rich_tracebacks=True,
show_time=False,
show_level=False,
show_path=False,
markup=True,
omit_repeated_times=False
)
rich_handler.setLevel(level)
rich_handler.setFormatter(self.formatter)
self.addHandler(rich_handler)
self.file_console = Console(
width=None,
markup=True,
color_system="truecolor",
force_terminal=True,
file=open(log_path, "a", encoding="utf-8")
)
rich_file_handler = RichHandler(
console=self.file_console,
rich_tracebacks=True,
show_time=False,
show_level=False,
show_path=False,
markup=True,
omit_repeated_times=False,
)
rich_file_handler.setLevel(level)
rich_file_handler.setFormatter(self.formatter)
self.addHandler(rich_file_handler)
self.propagate = False
def info(self, msg, *args, **kwargs):
"""
Only for string messages, not for rich objects.
"""
kwargs.setdefault("stacklevel", 2)
if "style" in kwargs:
kwargs.pop("style")
if "level" in kwargs:
kwargs.pop("level")
super().info(msg, *args, **kwargs)
def warning(self, msg, *args, **kwargs):
"""
Only for string messages, not for rich objects.
"""
kwargs.setdefault("stacklevel", 2)
super().warning(msg, *args, **kwargs)
def error(self, msg, *args, **kwargs):
kwargs.setdefault("stacklevel", 2)
super().error(msg, *args, **kwargs)
def critical(self, msg, *args, **kwargs):
kwargs.setdefault("stacklevel", 2)
super().critical(msg, *args, **kwargs)
def debug(self, msg, *args, **kwargs):
kwargs.setdefault("stacklevel", 2)
super().debug(msg, *args, **kwargs)
def log(self,
msg: Optional[Any] = None,
level: LogLevel = LogLevel.INFO,
**kwargs):
"""
Log a rich object or a string message to both console and file.
"""
if isinstance(msg, str):
self.info(msg, **kwargs)
elif isinstance(msg, (Group, Panel, Rule, Syntax, Table, Tree)):
self.console.print(msg, **kwargs)
self.file_console.print(msg, **kwargs)
logger = Logger() |