Make logs coloured!

1 files changed, 23 insertions, 7 deletions

diff --git a/bot/__init__.py b/bot/__init__.py
index f7a410706..c9dbc3f40 100644
--- a/bot/__init__.py
+++ b/bot/__init__.py
@@ -1,9 +1,11 @@
 import logging
 import os
 import sys
-from logging import Logger, StreamHandler, handlers
+from logging import Logger, handlers
 from pathlib import Path
 
+import coloredlogs
+
 TRACE_LEVEL = logging.TRACE = 5
 logging.addLevelName(TRACE_LEVEL, "TRACE")
 
@@ -25,10 +27,9 @@ Logger.trace = monkeypatch_trace
 
 DEBUG_MODE = 'local' in os.environ.get("SITE_URL", "local")
 
-log_format = logging.Formatter("%(asctime)s | %(name)s | %(levelname)s | %(message)s")
-
-stream_handler = StreamHandler(stream=sys.stdout)
-stream_handler.setFormatter(log_format)
+log_level = TRACE_LEVEL if DEBUG_MODE else logging.INFO
+format_string = "%(asctime)s | %(name)s | %(levelname)s | %(message)s"
+log_format = logging.Formatter(format_string)
 
 log_file = Path("logs", "bot.log")
 log_file.parent.mkdir(exist_ok=True)
@@ -36,10 +37,25 @@ file_handler = handlers.RotatingFileHandler(log_file, maxBytes=5242880, backupCo
 file_handler.setFormatter(log_format)
 
 root_log = logging.getLogger()
-root_log.setLevel(TRACE_LEVEL if DEBUG_MODE else logging.INFO)
-root_log.addHandler(stream_handler)
+root_log.setLevel(log_level)
 root_log.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
+
+if "COLOREDLOGS_LOG_LEVEL" not in os.environ:
+    coloredlogs.DEFAULT_LOG_LEVEL = log_level
+
+coloredlogs.install(logger=root_log, stream=sys.stdout)
+
 logging.getLogger("discord").setLevel(logging.WARNING)
 logging.getLogger("websockets").setLevel(logging.WARNING)
 logging.getLogger(__name__)