Python logger with colors and sensitive string masking

logger.py

#!/usr/bin/env python3
import logging
import os
from pathlib import Path

SENSITIVE: list[str] = []
TIMEZONE: str = "America/Denver"

# https://no-color.org/
NO_COLOR: bool = os.getenv("NO_COLOR", "") != ""

class Color:
    enabled = not NO_COLOR
    NC = "" if NO_COLOR else "\033[0m"  # No Color
    BLACK = "" if NO_COLOR else "\033[0;30m"
    GRAY = "" if NO_COLOR else "\033[1;30m"
    RED = "" if NO_COLOR else "\033[0;31m"
    LIGHT_RED = "" if NO_COLOR else "\033[1;31m"
    GREEN = "" if NO_COLOR else "\033[0;32m"
    LIGHT_GREEN = "" if NO_COLOR else "\033[1;32m"
    BROWN = "" if NO_COLOR else "\033[0;33m"
    YELLOW = "" if NO_COLOR else "\033[1;33m"
    BLUE = "" if NO_COLOR else "\033[0;34m"
    LIGHT_BLUE = "" if NO_COLOR else "\033[1;34m"
    PURPLE = "" if NO_COLOR else "\033[0;35m"
    LIGHT_PURPLE = "" if NO_COLOR else "\033[1;35m"
    CYAN = "" if NO_COLOR else "\033[0;36m"
    LIGHT_CYAN = "" if NO_COLOR else "\033[1;36m"
    LIGHT_GRAY = "" if NO_COLOR else "\033[0;37m"
    WHITE = "" if NO_COLOR else "\033[1;37m"

    @staticmethod
    def colorize(text: str, color) -> str:
        return f"{color}{text}{Color.NC}" if Color.enabled else text

class Formatter(logging.Formatter):
    color_levels = {
        logging.DEBUG: Color.GRAY,
        logging.INFO: Color.NC,
        logging.WARNING: Color.YELLOW,
        logging.ERROR: Color.RED,
        logging.CRITICAL: Color.RED,
    }

    def __init__(
        self,
        colors: bool = False,
        timezone: str = TIMEZONE,
        fmt="%(asctime)s %(name)s:%(lineno)03d [%(levelname)s] : %(message)s",
        *args,
        **kwargs
    ):
        self._colors = colors
        self._timezone = timezone

        super().__init__(
            fmt=fmt, datefmt="%m/%d/%Y %H:%M:%S %Z", *args, **kwargs
        )  # type: ignore

    def format(self, record):
        result = super().format(record)

        for item in SENSITIVE:
            result = result.replace(item, item[:2] + "*" * 8)

        if self._colors:
            result = self.color_levels.get(record.levelno, Color.NC) + result + Color.NC

        return result


def set_root_level(level):
    logging.getLogger().setLevel(level)


def init_logger(level: int = logging.INFO):
    handler = logging.StreamHandler()
    handler.setFormatter(Formatter(colors=True, fmt="%(message)s"))
    logging.basicConfig(level=level, handlers=[handler])


def add_log_file(filename: str | Path, level: int = logging.DEBUG):
    # Python logging is tricky.  The root logger has to be set to the lowest
    # level needed.  However, you don't neccessarily want to set the log level
    # for _everything_ (setting the root logger to `DEBUG` can set many library
    # loggers to `DEBUG` as well).
    root = logging.getLogger()
    if level < root.level:
        for handler in [x for x in root.handlers if not x.level]:
            handler.setLevel(root.level)

        root.setLevel(level)

    handler = logging.FileHandler(filename)
    handler.setFormatter(Formatter(colors=False))
    handler.setLevel(level)
    logging.getLogger().addHandler(handler)