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