1"""ANSI color support for MILC. 2""" 3import sys 4import re 5import logging 6import colorama 7 8from .emoji import EMOJI_LOGLEVELS 9 10ansi_config = { 11 'color': True, 12 'unicode': sys.stdout.encoding.lower().startswith('utf'), 13} 14 15# Regex was gratefully borrowed from kfir on stackoverflow: 16# https://stackoverflow.com/a/45448194 17ansi_regex = r'\x1b(' \ 18 r'(\[\??\d+[hl])|' \ 19 r'([=<>a-kzNM78])|' \ 20 r'([\(\)][a-b0-2])|' \ 21 r'(\[\d{0,2}[ma-dgkjqi])|' \ 22 r'(\[\d+;\d+[hfy]?)|' \ 23 r'(\[;?[hf])|' \ 24 r'(#[3-68])|' \ 25 r'([01356]n)|' \ 26 r'(O[mlnp-z]?)|' \ 27 r'(/Z)|' \ 28 r'(\d+)|' \ 29 r'(\[\?\d;\d0c)|' \ 30 r'(\d;\dR))' 31ansi_escape = re.compile(ansi_regex, flags=re.IGNORECASE) 32ansi_styles = ( 33 ('fg', colorama.ansi.AnsiFore()), 34 ('bg', colorama.ansi.AnsiBack()), 35 ('style', colorama.ansi.AnsiStyle()), 36) 37ansi_colors = {} 38 39for prefix, obj in ansi_styles: 40 for color in [x for x in obj.__dict__ if not x.startswith('_')]: 41 ansi_colors[prefix + '_' + color.lower()] = getattr(obj, color) 42 43 44def format_ansi(text): 45 """Return a copy of text with certain strings replaced with ansi. 46 """ 47 # Avoid .format() so we don't have to worry about the log content 48 for color in ansi_colors: 49 text = text.replace('{%s}' % color, ansi_colors[color]) 50 51 text = text + ansi_colors['style_reset_all'] 52 53 if ansi_config['color']: 54 return text 55 56 return ansi_escape.sub('', text) 57 58 59class MILCFormatter(logging.Formatter): 60 """Formats log records per the MILC configuration. 61 """ 62 def format(self, record): 63 if ansi_config['unicode'] and record.levelname in EMOJI_LOGLEVELS: 64 record.levelname = format_ansi(EMOJI_LOGLEVELS[record.levelname]) 65 66 msg = super().format(record) 67 return format_ansi(msg) 68