Add `CustomLogger` to bot/log.py

1 files changed, 33 insertions, 18 deletions

diff --git a/bot/log.py b/bot/log.py
index 4e20c005e..74d8263c0 100644
--- a/bot/log.py
+++ b/bot/log.py
@@ -3,6 +3,7 @@ import os
 import sys
 from logging import Logger, handlers
 from pathlib import Path
+from typing import Optional, TYPE_CHECKING, cast
 
 import coloredlogs
 import sentry_sdk
@@ -14,11 +15,38 @@ from bot import constants
 TRACE_LEVEL = 5
 
 
+if TYPE_CHECKING:
+    LoggerClass = Logger
+else:
+    LoggerClass = logging.getLoggerClass()
+
+
+class CustomLogger(LoggerClass):
+    """Custom implementation of the `Logger` class with an added `trace` method."""
+
+    def trace(self, msg: str, *args, **kwargs) -> None:
+        """
+        Log 'msg % args' with severity 'TRACE'.
+
+        To pass exception information, use the keyword argument exc_info with
+        a true value, e.g.
+
+        logger.trace("Houston, we have an %s", "interesting problem", exc_info=1)
+        """
+        if self.isEnabledFor(TRACE_LEVEL):
+            self.log(TRACE_LEVEL, msg, *args, **kwargs)
+
+
+def get_logger(name: Optional[str] = None) -> CustomLogger:
+    """Utility to make mypy recognise that logger is of type `CustomLogger`."""
+    return cast(CustomLogger, logging.getLogger(name))
+
+
 def setup() -> None:
     """Set up loggers."""
     logging.TRACE = TRACE_LEVEL
     logging.addLevelName(TRACE_LEVEL, "TRACE")
-    Logger.trace = _monkeypatch_trace
+    logging.setLoggerClass(CustomLogger)
 
     format_string = "%(asctime)s | %(name)s | %(levelname)s | %(message)s"
     log_format = logging.Formatter(format_string)
@@ -42,7 +70,7 @@ def setup() -> None:
     if "COLOREDLOGS_LOG_FORMAT" not in os.environ:
         coloredlogs.DEFAULT_LOG_FORMAT = format_string
 
-    coloredlogs.install(level=logging.TRACE, logger=root_log, stream=sys.stdout)
+    coloredlogs.install(level=TRACE_LEVEL, logger=root_log, stream=sys.stdout)
 
     root_log.setLevel(logging.DEBUG if constants.DEBUG_MODE else logging.INFO)
     logging.getLogger("discord").setLevel(logging.WARNING)
@@ -73,19 +101,6 @@ def setup_sentry() -> None:
     )
 
 
-def _monkeypatch_trace(self: logging.Logger, msg: str, *args, **kwargs) -> None:
-    """
-    Log 'msg % args' with severity 'TRACE'.
-
-    To pass exception information, use the keyword argument exc_info with
-    a true value, e.g.
-
-    logger.trace("Houston, we have an %s", "interesting problem", exc_info=1)
-    """
-    if self.isEnabledFor(TRACE_LEVEL):
-        self._log(TRACE_LEVEL, msg, args, **kwargs)
-
-
 def _set_trace_loggers() -> None:
     """
     Set loggers to the trace level according to the value from the BOT_TRACE_LOGGERS env var.
@@ -101,13 +116,13 @@ def _set_trace_loggers() -> None:
     level_filter = constants.Bot.trace_loggers
     if level_filter:
         if level_filter.startswith("*"):
-            logging.getLogger().setLevel(logging.TRACE)
+            logging.getLogger().setLevel(TRACE_LEVEL)
 
         elif level_filter.startswith("!"):
-            logging.getLogger().setLevel(logging.TRACE)
+            logging.getLogger().setLevel(TRACE_LEVEL)
             for logger_name in level_filter.strip("!,").split(","):
                 logging.getLogger(logger_name).setLevel(logging.DEBUG)
 
         else:
             for logger_name in level_filter.strip(",").split(","):
-                logging.getLogger(logger_name).setLevel(logging.TRACE)
+                logging.getLogger(logger_name).setLevel(TRACE_LEVEL)