Merge pull request #858 from Numerlor/coloredlogs

4 files changed, 99 insertions, 60 deletions

diff --git a/bot/__init__.py b/bot/__init__.py
index cfaee9f8..b19bd76a 100644
--- a/bot/__init__.py
+++ b/bot/__init__.py
@@ -6,66 +6,19 @@ except ModuleNotFoundError:
     pass
 
 import asyncio
-import logging
-import logging.handlers
 import os
 from functools import partial, partialmethod
-from pathlib import Path
 
 import arrow
 from discord.ext import commands
 
-from bot import monkey_patches
-from bot.constants import Client
+from bot import log, monkey_patches
 
-# Configure the "TRACE" logging level (e.g. "log.trace(message)")
-logging.TRACE = 5
-logging.addLevelName(logging.TRACE, "TRACE")
-
-logging.Logger.trace = monkey_patches.trace_log
+log.setup()
 
 # Set timestamp of when execution started (approximately)
 start_time = arrow.utcnow()
 
-# Set up file logging
-log_dir = Path("bot/log")
-log_file = log_dir / "hackbot.log"
-os.makedirs(log_dir, exist_ok=True)
-
-# File handler rotates logs every 5 MB
-file_handler = logging.handlers.RotatingFileHandler(
-    log_file, maxBytes=5 * (2**20), backupCount=10, encoding="utf-8",
-)
-file_handler.setLevel(logging.TRACE if Client.debug else logging.DEBUG)
-
-# Console handler prints to terminal
-console_handler = logging.StreamHandler()
-level = logging.TRACE if Client.debug else logging.INFO
-console_handler.setLevel(level)
-
-# Remove old loggers, if any
-root = logging.getLogger()
-if root.handlers:
-    for handler in root.handlers:
-        root.removeHandler(handler)
-
-# Silence irrelevant loggers
-logging.getLogger("discord").setLevel(logging.ERROR)
-logging.getLogger("websockets").setLevel(logging.ERROR)
-logging.getLogger("PIL").setLevel(logging.ERROR)
-logging.getLogger("matplotlib").setLevel(logging.ERROR)
-logging.getLogger("async_rediscache").setLevel(logging.WARNING)
-
-# Setup new logging configuration
-logging.basicConfig(
-    format="%(asctime)s - %(name)s %(levelname)s: %(message)s",
-    datefmt="%D %H:%M:%S",
-    level=logging.TRACE if Client.debug else logging.DEBUG,
-    handlers=[console_handler, file_handler],
-)
-logging.getLogger().info("Logging initialization complete")
-
-
 # On Windows, the selector event loop is required for aiodns.
 if os.name == "nt":
     asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
diff --git a/bot/constants.py b/bot/constants.py
index 0720dd20..9d12000e 100644
--- a/bot/constants.py
+++ b/bot/constants.py
@@ -139,6 +139,7 @@ class Client(NamedTuple):
     github_bot_repo = "https://github.com/python-discord/sir-lancebot"
     # Override seasonal locks: 1 (January) to 12 (December)
     month_override = int(environ["MONTH_OVERRIDE"]) if "MONTH_OVERRIDE" in environ else None
+    trace_loggers = environ.get("BOT_TRACE_LOGGERS")
 
 
 class Colours:
diff --git a/bot/log.py b/bot/log.py
new file mode 100644
index 00000000..5e0e909d
--- /dev/null
+++ b/bot/log.py
@@ -0,0 +1,96 @@
+import logging
+import logging.handlers
+import os
+import sys
+from pathlib import Path
+
+import coloredlogs
+
+from bot.constants import Client
+
+
+def setup() -> None:
+    """Set up loggers."""
+    # Configure the "TRACE" logging level (e.g. "log.trace(message)")
+    logging.TRACE = 5
+    logging.addLevelName(logging.TRACE, "TRACE")
+    logging.Logger.trace = _monkeypatch_trace
+
+    format_string = "%(asctime)s | %(name)s | %(levelname)s | %(message)s"
+    log_format = logging.Formatter(format_string)
+
+    # Set up file logging
+    log_file = Path("logs/sir-lancebot.log")
+    log_file.parent.mkdir(exist_ok=True)
+
+    # File handler rotates logs every 5 MB
+    file_handler = logging.handlers.RotatingFileHandler(
+        log_file, maxBytes=5 * (2 ** 20), backupCount=10, encoding="utf-8",
+    )
+    file_handler.setFormatter(log_format)
+
+    root_logger = logging.getLogger()
+    root_logger.addHandler(file_handler)
+
+    if "COLOREDLOGS_LEVEL_STYLES" not in os.environ:
+        coloredlogs.DEFAULT_LEVEL_STYLES = {
+            **coloredlogs.DEFAULT_LEVEL_STYLES,
+            "trace": {"color": 246},
+            "critical": {"background": "red"},
+            "debug": coloredlogs.DEFAULT_LEVEL_STYLES["info"],
+        }
+
+    if "COLOREDLOGS_LOG_FORMAT" not in os.environ:
+        coloredlogs.DEFAULT_LOG_FORMAT = format_string
+
+    coloredlogs.install(level=logging.TRACE, stream=sys.stdout)
+
+    root_logger.setLevel(logging.DEBUG if Client.debug else logging.INFO)
+    # Silence irrelevant loggers
+    logging.getLogger("discord").setLevel(logging.ERROR)
+    logging.getLogger("websockets").setLevel(logging.ERROR)
+    logging.getLogger("PIL").setLevel(logging.ERROR)
+    logging.getLogger("matplotlib").setLevel(logging.ERROR)
+    logging.getLogger("async_rediscache").setLevel(logging.WARNING)
+
+    _set_trace_loggers()
+
+    root_logger.info("Logging initialization complete")
+
+
+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(logging.TRACE):
+        self._log(logging.TRACE, 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.
+
+    When the env var is a list of logger names delimited by a comma,
+    each of the listed loggers will be set to the trace level.
+
+    If this list is prefixed with a "!", all of the loggers except the listed ones will be set to the trace level.
+
+    Otherwise if the env var begins with a "*",
+    the root logger is set to the trace level and other contents are ignored.
+    """
+    level_filter = Client.trace_loggers
+    if level_filter:
+        if level_filter.startswith("*"):
+            logging.getLogger().setLevel(logging.TRACE)
+
+        elif level_filter.startswith("!"):
+            logging.getLogger().setLevel(logging.TRACE)
+            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)
diff --git a/bot/monkey_patches.py b/bot/monkey_patches.py
index fe81f2e3..fa6627d1 100644
--- a/bot/monkey_patches.py
+++ b/bot/monkey_patches.py
@@ -7,17 +7,6 @@ from discord.ext import commands
 log = logging.getLogger(__name__)
 
 
-def trace_log(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(logging.TRACE):
-        self._log(logging.TRACE, msg, args, **kwargs)
-
-
 class Command(commands.Command):
     """
     A `discord.ext.commands.Command` subclass which supports root aliases.