Source code for jgdv.logging.format.colour

  1#!/usr/bin/env python3
  2"""
  3see `Alexandra Zaharia <https://alexandra-zaharia.github.io/posts/make-your-own-custom-color-formatter-with-python-logging/>`_
  4
  5"""
  6# Imports:
  7from __future__ import annotations
  8
  9# ##-- stdlib imports
 10import datetime
 11import enum
 12import functools as ftz
 13import itertools as itz
 14import logging
 15import logging as logmod
 16import pathlib as pl
 17import re
 18import warnings
 19from collections import defaultdict
 20from string import Formatter
 21from uuid import UUID, uuid1
 22
 23# ##-- end stdlib imports
 24
 25# ##-- 3rd party imports
 26import sty
 27from sty import bg, ef, fg, rs
 28
 29# ##-- end 3rd party imports
 30
 31# ##-- 1st party imports
 32from jgdv import Mixin, Proto
 33
 34# ##-- end 1st party imports
 35
 36from . import _interface as API # noqa: N812
 37from .stack_m import StackFormatter_m
 38
 39# ##-- types
 40# isort: off
 41import abc
 42import collections.abc
 43from typing import TYPE_CHECKING, cast, assert_type, assert_never
 44from typing import Generic, NewType
 45# Protocols:
 46from typing import Protocol, runtime_checkable
 47# Typing Decorators:
 48from typing import no_type_check, final, override, overload
 49
 50if TYPE_CHECKING:
 51    from jgdv import Maybe, Rx
 52    from typing import Final
 53    from typing import ClassVar, Any, LiteralString
 54    from typing import Never, Self, Literal
 55    from typing import TypeGuard
 56    from collections.abc import Iterable, Iterator, Callable, Generator
 57    from collections.abc import Sequence, Mapping, MutableMapping, Hashable
 58
 59    from logging import LogRecord
 60
 61    type StyleChar = Literal["%", "{", "$"]
 62##--|
 63
 64# isort: on
 65# ##-- end types
 66
 67COLOUR_RESET       : str    = rs.all
 68##--|
 69
 70

[docs]
 71@Mixin(StackFormatter_m)
 72class ColourFormatter(logging.Formatter):
 73    """
 74    Stream Formatter for logging, enables use of colour sent to console
 75
 76    Guarded Formatter for adding colour.
 77    Uses the sty module.
 78    If sty is missing, behaves as the default formatter class
 79
 80    # Do *not* use for on filehandler
 81    Usage reminder:
 82    # Create stdout handler for logging to the console (logs all five levels)
 83    stdout_handler = logging.StreamHandler()
 84    stdout_handler.setFormatter(ColourFormatter(fmt))
 85    logger.addHandler(stdout_handler)
 86    """
 87
 88    _default_fmt      : str        = '{asctime} | {levelname:9} | {message}'
 89    _default_date_fmt : str        =  "%H:%M:%S"
 90    _default_style    : StyleChar  = '{'
 91    colours           : dict[int|str, str]
 92
 93    def __init__(self, *, fmt:Maybe[str]=None, style:Maybe[StyleChar]=None) -> None:
 94        """
 95        Create the ColourFormatter with a given *Brace* style log format
 96        """
 97        super().__init__(fmt or self._default_fmt,
 98                         datefmt=self._default_date_fmt,
 99                         style=style or self._default_style)
100        self.colours = defaultdict(lambda: rs.all)
101        self.apply_colour_mapping(API.default_colour_mapping)
102        self.apply_colour_mapping(API.default_log_colours)
103

[docs]
104    @override
105    def format(self, record:LogRecord) -> str:
106        if hasattr(record, "colour"):
107            log_colour = self.colours[record.colour]
108        else:
109            log_colour = self.colours[record.levelno]
110
111        return log_colour + super().format(record) + COLOUR_RESET

112

[docs]
113    def apply_colour_mapping(self, mapping:dict) -> None:
114        """ applies a mapping of colours by treating each value as a pair of attrs of sty
115
116        eg: {logging.DEBUG: ("fg", "blue"), logging.INFO: ("bg", "red")}
117        """
118        for x,(a,b) in mapping.items():
119            accessor = getattr(sty, a)
120            val      = getattr(accessor, b)
121            self.colours[x] = val


122
123

[docs]
124@Mixin(StackFormatter_m)
125class StripColourFormatter(logging.Formatter):
126    """
127    Force Colour Command codes to be stripped out of a string.
128    Useful for when you redirect printed strings with colour
129    to a file
130    """
131
132    _default_fmt      : str           = "{asctime} | {levelname:9} | {shortname:25} | {message}"
133    _default_date_fmt : str           = "%Y-%m-%d %H:%M:%S"
134    _default_style    : StyleChar     = '{'
135    _colour_strip_re  : Rx            = re.compile(r'\x1b\[([\d;]+)m?')
136
137    def __init__(self, *, fmt:Maybe[str]=None, style:Maybe[StyleChar]=None) -> None:
138        """
139        Create the StripColourFormatter with a given *Brace* style log format
140        """
141        super().__init__(fmt or self._default_fmt,
142                         datefmt=style or self._default_date_fmt,
143                         style=self._default_style)
144

[docs]
145    @override
146    def format(self, record:LogRecord) -> str:
147        result    = super().format(record)
148        no_colour = self._colour_strip_re.sub("", result)
149        return no_colour