Merge pull request #2505 from cta-observatory/fix_unknown_logging

Improve logging colorization, fixes #2504

docs/changes/2505.bugfix.rst

@@ -0,0 +1 @@
+Fix colored logging in case of custom log levels being defined.

src/ctapipe/core/logging.py

@@ -2,12 +2,44 @@
 
 import logging
 from collections.abc import Mapping
+from enum import IntEnum
 
 DEFAULT_LOGGING_FORMAT = (
     "%(asctime)s %(levelname)s [%(name)s] (%(module)s.%(funcName)s): %(message)s"
 )
 
 
+class ANSIEscapes(IntEnum):
+    """
+    Enum of some ANSI style escape sequences used for logging.
+
+    See https://en.wikipedia.org/wiki/ANSI_escape_code
+    """
+
+    RESET = 0
+    BOLD = 1
+    FG_RED = 31
+    FG_GREEN = 32
+    FG_YELLOW = 33
+    FG_BLUE = 34
+    FG_MAGENTA = 35
+
+
+LEVEL_COLORS = {
+    "DEBUG": ANSIEscapes.FG_BLUE,
+    "INFO": ANSIEscapes.FG_GREEN,
+    "WARNING": ANSIEscapes.FG_YELLOW,
+    "ERROR": ANSIEscapes.FG_RED,
+    "CRITICAL": ANSIEscapes.FG_MAGENTA,
+}
+
+
+def add_ansi_display(string, *styles):
+    """Surround ``string`` by ANSI escape code styling."""
+    styles = ";".join(str(int(style)) for style in styles)
+    return f"\033[{styles}m{string}\033[{ANSIEscapes.RESET}m"
+
+
 class PlainFormatter(logging.Formatter):
     """Custom logging.Formatter used for file logging."""
 
@@ -23,20 +55,12 @@ def format(self, record):
 
 def apply_colors(levelname: str):
     """Use ANSI escape sequences to add colors the levelname of log entries."""
-    _black, red, green, yellow, blue, magenta, _cyan, _white = range(8)
-    reset_seq = "\033[0m"
-    color_seq = "\033[1;%dm"
-    colors = {
-        "INFO": green,
-        "DEBUG": blue,
-        "WARNING": yellow,
-        "CRITICAL": magenta,
-        "ERROR": red,
-    }
+    color = LEVEL_COLORS.get(levelname)
 
-    levelname_color = color_seq % (30 + colors[levelname]) + levelname + reset_seq
+    if color is None:
+        return add_ansi_display(levelname, ANSIEscapes.BOLD)
 
-    return levelname_color
+    return add_ansi_display(levelname, ANSIEscapes.BOLD, color)
 
 
 def recursive_update(d1, d2, copy=False):

src/ctapipe/core/tests/test_logging.py

@@ -0,0 +1,12 @@
+def test_colors():
+    """Test applying colors"""
+    from ctapipe.core.logging import apply_colors
+
+    assert apply_colors("DEBUG") == "\x1b[1;34mDEBUG\x1b[0m"
+    assert apply_colors("INFO") == "\x1b[1;32mINFO\x1b[0m"
+    assert apply_colors("WARNING") == "\x1b[1;33mWARNING\x1b[0m"
+    assert apply_colors("ERROR") == "\x1b[1;31mERROR\x1b[0m"
+    assert apply_colors("CRITICAL") == "\x1b[1;35mCRITICAL\x1b[0m"
+
+    # unknown log-level, regression test for #2504
+    assert apply_colors("FOO") == "\x1b[1mFOO\x1b[0m"