Allow changing logger name

youtype · Jan 17, 2025 · 0923876 · 0923876
1 parent 8d532dc
commit 0923876
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/mypy_boto3_builder/logger.py b/mypy_boto3_builder/logger.py
@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import functools
 import logging
 import sys
 
@@ -16,27 +17,29 @@
 __all__ = ("get_logger", "setup_logger")
 
 
-def _formatter(record: loguru.Record) -> str:
+def _formatter(name: str, record: loguru.Record) -> str:
     tags = record["extra"].get("tags") or ()
     message = record["message"]
     for tag in tags:
         message = message.replace(f"{tag}", f"<cyan>{tag}</cyan>")
     return (
         "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
         "<level>{level: <8}</level> | "
-        f"<cyan>{LOGGER_NAME}</cyan> - "
+        f"<cyan>{name}</cyan> - "
         f"<white>{message}</white>"
         "\n"
     )
 
 
-def setup_logger(level: int) -> None:
+def setup_logger(level: int = logging.DEBUG, name: str = LOGGER_NAME) -> None:
     """
     Set up logger.
     """
     level_name = logging.getLevelName(level)
     loguru.logger.configure(
-        handlers=[{"sink": sys.stderr, "level": level_name, "format": _formatter}]
+        handlers=[
+            {"sink": sys.stderr, "level": level_name, "format": functools.partial(_formatter, name)}
+        ],
     )