make a colorize module, refactor colored logging to use it

Signed-off-by: Zen <[email protected]>

Author: desultory

src/zenlib/__init__.py

@@ -1,8 +1,9 @@
 from .logging import ColorLognameFormatter, loggify
 from .types import NoDupFlatList, validatedDataclass
-from .util import check_dict, handle_plural, pretty_print, replace_file_line, update_init, walk_dict
+from .util import colorize, check_dict, handle_plural, pretty_print, replace_file_line, update_init, walk_dict
 
 __all__ = [
+    "colorize",
     "ColorLognameFormatter",
     "loggify",
     "handle_plural",

src/zenlib/logging/colorlognameformatter.py

@@ -1,47 +1,34 @@
 __author__ = "desultory"
-__version__ = "2.0.0"
-
+__version__ = "3.0.0"
 
 from logging import Formatter
 
+from zenlib.util import colorize
 
-class ColorLognameFormatter(Formatter):
-    """
-    ColorLognameFormatter Class
-    Add the handler to the stdout handler using:
-    stdout_handler = logging.StreamHandler()
-    stdout_handler.setFormatter(ColorLognameFormatter())
-    """
-
-    # Define the color codes
-    _reset_str = '\x1b[0m'
-    _grey_str = '\x1b[37m'
-    _blue_str = '\x1b[34m'
-    _yllw_str = '\x1b[33m'
-    _sred_str = '\x1b[31m'
-    _bred_str = '\x1b[31;1m'
-    _magenta_str = '\x1b[35m'
-    # Format into a dict
-    _color_levelname = {'DEBUG': _grey_str,
-                        'INFO': _blue_str,
-                        'WARNING': _yllw_str,
-                        'ERROR': _sred_str,
-                        'CRITICAL': _bred_str}
 
-    def __init__(self, fmt='%(levelname)s | %(message)s', *args, **kwargs):
+class ColorLognameFormatter(Formatter):
+    """A logging formatter which colors the levelname of the log message.
+    Uses the zenlib.util.colorize function to color the levelname.
+    Normal levelnames are padded to the length of the longest levelname."""
+
+    level_colors = {
+        "DEBUG": {"color": "white"},
+        "INFO": {"color": "blue"},
+        "WARNING": {"color": "yellow"},
+        "ERROR": {"color": "red", "bold": True},
+        "CRITICAL": {"color": "red", "bold": True, "bright": True},
+    }
+
+    def __init__(self, fmt="%(levelname)s | %(message)s", *args, **kwargs):
         super().__init__(fmt, *args, **kwargs)
-        self._level_str_len = max([len(name) for name in self._color_levelname])
-
-    def color_levelname(self, levelname):
-        """ Adds color to a levelname string """
-        color = self._color_levelname.get(levelname, self._magenta_str)
-        return f"{color}{levelname}{self._reset_str}".ljust(self._level_str_len + len(color) + len(self._reset_str), ' ')
+        self.level_str_width = max(len(name) for name in self.level_colors) - 1
 
     def format(self, record):
         # When calling format, replace the levelname with a colored version
         # Note: the string size is greatly increased because of the color codes
         old_levelname = record.levelname
-        record.levelname = self.color_levelname(record.levelname)
+        color_info = self.level_colors.get(record.levelname, {"color": "white"})
+        record.levelname = colorize(record.levelname.ljust(self.level_str_width), **color_info)
         format_str = super().format(record)
         record.levelname = old_levelname
         return format_str

src/zenlib/util/__init__.py

@@ -1,4 +1,5 @@
 from .dict_check import contains, unset
+from .colorize import colorize
 from .handle_plural import handle_plural
 from .main_funcs import get_args_n_logger, get_kwargs, get_kwargs_from_args, init_argparser, init_logger, process_args
 from .merge_class import merge_class
@@ -7,6 +8,7 @@
 
 __all__ = [
     "handle_plural",
+    "colorize",
     "NoDupFlatList",
     "pretty_print",
     "replace_file_line",

src/zenlib/util/colorize.py

@@ -0,0 +1,33 @@
+__version__ = "0.1.0"
+
+from enum import Enum
+
+
+ANSI_START = "\x1b["
+
+class Colors(Enum):
+    RED = 31
+    GREEN = 32
+    YELLOW = 33
+    BLUE = 34
+    MAGENTA = 35
+    CYAN = 36
+    WHITE = 37
+
+
+class ANSICode:
+    def __init__(self, code, bright=False, bold=False):
+        self.code = code
+        self.bright = bright
+        self.bold = bold
+
+    def __str__(self):
+        color_code = self.code
+        color_code += 60 if self.bright else 0
+        bold_str = ";1" if self.bold else ""
+        return f"{ANSI_START}{color_code}{bold_str}m"
+
+
+def colorize(text: str, color: str, bright=False, bold=False) -> str:
+    color_code = Colors[color.upper()].value
+    return f"{ANSICode(color_code, bright, bold)}{text}{ANSICode(0)}"