feat: add log RichHandler & RotatingFileHandler

Author: ChenZiHong-Gavin

graphgen/utils/log.py

@@ -1,32 +1,45 @@
 import logging
+from logging.handlers import RotatingFileHandler
 
-logger = logging.getLogger("graphgen")
+from rich.logging import RichHandler
 
-def set_logger(log_file: str, log_level: int = logging.INFO, if_stream: bool = True):
-    logger.setLevel(log_level)
+logger = logging.getLogger("graphgen")
 
-    formatter = logging.Formatter(
-        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-    )
 
-    file_handler = logging.FileHandler(log_file, mode='w')
-    file_handler.setLevel(log_level)
-    file_handler.setFormatter(formatter)
+def set_logger(
+    log_file: str,
+    log_level: int = logging.INFO,
+    if_stream: bool = True,
+    max_bytes: int = 50 * 1024 * 1024,  # 50 MB
+    backup_count: int = 5,
+):
+    logger.setLevel(log_level)
+    logger.propagate = False
 
-    stream_handler = None
+    if logger.handlers:
+        logger.handlers.clear()
 
     if if_stream:
-        stream_handler = logging.StreamHandler()
-        stream_handler.setLevel(log_level)
-        stream_handler.setFormatter(formatter)
-
-    if not logger.handlers:
-        logger.addHandler(file_handler)
-        if if_stream and stream_handler:
-            logger.addHandler(stream_handler)
+        console = RichHandler(level=log_level, show_path=False, rich_tracebacks=True)
+        console.setFormatter(logging.Formatter("%(message)s"))
+        logger.addHandler(console)
+
+    file_handler = RotatingFileHandler(
+        log_file,
+        maxBytes=max_bytes,
+        backupCount=backup_count,
+        encoding="utf-8",
+    )
+    file_handler.setFormatter(
+        logging.Formatter(
+            "[%(asctime)s] %(levelname)s [%(name)s:%(filename)s:%(lineno)d] %(message)s",
+            datefmt="%y-%m-%d %H:%M:%S",
+        )
+    )
+    logger.addHandler(file_handler)
 
 
 def parse_log(log_file: str):
-    with open(log_file, "r", encoding='utf-8') as f:
+    with open(log_file, "r", encoding="utf-8") as f:
         lines = f.readlines()
     return lines