Optimize logger

src/common/utils/logger.py

@@ -5,39 +5,13 @@ from logging import (
     INFO,
     WARNING,
     Logger,
-    _levelToName,
     _nameToLevel,
     addLevelName,
     basicConfig,
     getLogger,
-    setLoggerClass,
 )
 from os import getenv
-from threading import Lock
-
-
-class BWLogger(Logger):
-    def __init__(self, name, level=INFO):
-        self.name = name
-        self.db_lock = Lock()
-        return super(BWLogger, self).__init__(name, level)
-
-    def _log(
-        self,
-        level,
-        msg,
-        args,
-        exc_info=None,
-        extra=None,
-        stack_info=False,
-        stacklevel=1,
-    ):
-        return super(BWLogger, self)._log(
-            level, msg, args, exc_info, extra, stack_info, stacklevel
-        )
-
-
-setLoggerClass(BWLogger)
+from typing import Optional, Union
 
 default_level = _nameToLevel.get(getenv("LOG_LEVEL", "INFO").upper(), INFO)
 basicConfig(
@@ -73,10 +47,15 @@ addLevelName(INFO, "ℹ️ ")
 addLevelName(WARNING, "⚠️ ")
 
 
-def setup_logger(title: str, level=INFO) -> Logger:
+def setup_logger(title: str, level: Optional[Union[str, int]] = None) -> Logger:
     """Set up local logger"""
     title = title.upper()
     logger = getLogger(title)
-    logger.setLevel(_nameToLevel.get(level, _levelToName.get(level, INFO)))
+    level = level or default_level
+
+    if isinstance(level, str):
+        logger.setLevel(_nameToLevel.get(level.upper(), default_level))
+    else:
+        logger.setLevel(level)
 
     return logger