Add code for parsing ANSI escape codes into Urwid AttrSpec.

Based on code by Patrick Totzke <patricktotzke@gmail.com>.

1 files changed, 216 insertions, 0 deletions

diff --git a/alot/utils/ansi_term.py b/alot/utils/ansi_term.py
new file mode 100644
index 00000000..291414a6
--- /dev/null
+++ b/alot/utils/ansi_term.py
@@ -0,0 +1,216 @@
+# Copyright (C) 2011-2017  Patrick Totzke <patricktotzke@gmail.com>
+# Copyright (C) 2021  Anton Khirnov <anton@khirnov.net>
+# This file is released under the GNU GPL, version 3 or a later revision.
+# For further details see the COPYING file
+
+import logging
+import re
+
+import urwid
+
+# see https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_sequences
+_SGR_CODES = {
+    # attributes ON
+    '1'     : { 'bold'          : True },
+    '3'     : { 'italics'       : True },
+    '4'     : { 'underline'     : True },
+    '5'     : { 'blink'         : True },
+    '7'     : { 'standout'      : True },
+    '9'     : { 'strikethrough' : True },
+
+    # attributes OFF
+    '22'    : { 'bold'          : False },
+    '23'    : { 'italics'       : False },
+    '24'    : { 'underline'     : False },
+    '25'    : { 'blink'         : False },
+    '27'    : { 'standout'      : False },
+    '29'    : { 'strikethrough' : False },
+
+    # foreground color
+    '30'    : { 'fg'            : 'black'},
+    '31'    : { 'fg'            : 'dark red'},
+    '32'    : { 'fg'            : 'dark green'},
+    '33'    : { 'fg'            : 'brown'},
+    '34'    : { 'fg'            : 'dark blue'},
+    '35'    : { 'fg'            : 'dark magenta'},
+    '36'    : { 'fg'            : 'dark cyan'},
+    '37'    : { 'fg'            : 'light gray'},
+
+    # background color
+    '40'    : { 'bg'            : 'black'},
+    '41'    : { 'bg'            : 'dark red'},
+    '42'    : { 'bg'            : 'dark green'},
+    '43'    : { 'bg'            : 'brown'},
+    '44'    : { 'bg'            : 'dark blue'},
+    '45'    : { 'bg'            : 'dark magenta'},
+    '46'    : { 'bg'            : 'dark cyan'},
+    '47'    : { 'bg'            : 'light gray'},
+
+    # bright foreground color
+    '90'    : { 'fg'            : 'dark gray'},
+    '91'    : { 'fg'            : 'light red'},
+    '92'    : { 'fg'            : 'light green'},
+    '93'    : { 'fg'            : 'yellow'},
+    '94'    : { 'fg'            : 'light blue'},
+    '95'    : { 'fg'            : 'light magenta'},
+    '96'    : { 'fg'            : 'light cyan'},
+    '97'    : { 'fg'            : 'white'},
+
+    # bright background color
+    '100'   : { 'bg'            : 'dark gray'},
+    '101'   : { 'bg'            : 'light red'},
+    '102'   : { 'bg'            : 'light green'},
+    '103'   : { 'bg'            : 'yellow'},
+    '104'   : { 'bg'            : 'light blue'},
+    '105'   : { 'bg'            : 'light magenta'},
+    '106'   : { 'bg'            : 'light cyan'},
+    '107'   : { 'bg'            : 'white'},
+}
+
+class _TermState(dict):
+    _ATTRS = (
+        'bold',
+        'underline',
+        'standout',
+        'blink',
+        'italics',
+        'strikethrough',
+    )
+
+    _default = None
+
+    def __init__(self, default_attr):
+        super().__init__()
+
+        self._default = default_attr
+
+        self._reset_attr()
+        self._reset_fg()
+        self._reset_bg()
+
+    def _reset_fg(self):
+        self['fg'] = self._default.foreground
+    def _reset_bg(self):
+        self['bg'] = self._default.background
+    def _reset_attr(self):
+        for attr in self._ATTRS:
+            self[attr] = getattr(self._default, attr)
+
+    def process_sgr_escape(self, pb, ib, fb):
+        if fb != 'm':
+            return
+
+        if not pb:
+            pb = '0'
+
+        while pb:
+            code, _, pb = pb.partition(';')
+
+            if code in _SGR_CODES:
+                self.update(_SGR_CODES[code])
+            elif code == '0':
+                self._reset_attr()
+            elif code == '39':
+                self._reset_fg()
+            elif code == '49':
+                self._reset_bg()
+            elif code == '38' or code == '48':
+                dst   = 'fg' if code == '38' else 'bg'
+                color = None
+
+                mode, _, pb = pb.partition(';')
+                if mode == '5':
+                    # 8-bit color index
+                    val, _, pb = pb.partition(';')
+                    try:
+                        val = int(val)
+                        if val >= 0 and val <= 255:
+                            color = 'h%d' % val
+                    except ValueError:
+                        logging.warning('Invalid 8-bit color index: %s', val)
+                elif mode == '2':
+                    # 24-bit RGB
+                    r, _, pb = pb.partition(';')
+                    g, _, pb = pb.partition(';')
+                    b, _, pb = pb.partition(';')
+
+                    try:
+                        r = int(r)
+                        g = int(g)
+                        b = int(b)
+                        if (r >= 0 and r <= 255 and g >= 0 and g <= 255 and
+                            b >= 0 and b <= 255):
+                            color = '#%x%x%x' % (r, g, b)
+                    except ValueError:
+                        logging.warning('Invalid 24-bit color specification: %s %s %s', r, g, b)
+
+                if color is not None:
+                    self[dst] = color
+
+
+    def to_urwid(self, preserve_bg):
+        fg = self['fg']
+        for attr in self._ATTRS:
+            if self[attr]:
+                urwid_fg += ',' + attr
+        bg = self._default.background if preserve_bg else self['bg']
+
+        return urwid.AttrSpec(fg, bg)
+
+def parse_escapes_to_urwid(text, default_attr = None, preserve_bg = False):
+    """
+    Process text containing ANSI CSI escape codes. From those, the SGR (Select
+    Graphic Rendition) sequences are converted to Urwid AttrSpec objects and
+    used to style corresponding parts of text. All CSI escape sequences are
+    removed from text, other escape codes are unaffected.
+
+    :param str text the text to process
+    :param urwid.AttrSpec default_attr default style to apply when no escape
+        code is in effect
+    :param preserve_bg bool preserve original background, i.e. disregard escape
+        codes affecting background
+    :return a list of parsed text blocks, where each block is a tuple of
+        (attribute, text).
+    :rtype list of (urwid.AttrSpec, str)
+    """
+
+    b1 = r'\033\['                  # Control Sequence Introducer
+    b2 = r'[0-9:;<=>?]*'            # parameter bytes
+    b3 = r'[ !\"#$%&\'()*+,-./]*'   # intermediate bytes
+    b4 = r'[A-Z[\]^_`a-z{|}~]'      # final byte"
+    esc_pattern = b1 \
+        + r'(?P<pb>' + b2 + ')' \
+        + r'(?P<ib>' + b3 + ')' \
+        + r'(?P<fb>' + b4 + ')'
+
+    state = _TermState(default_attr)
+    ret   = []
+
+    def append_themed_block(block):
+        lines = block.splitlines(keepends = True)
+
+        # since blocks are terminated on escape codes and some formatters reset
+        # attributes right before the end of each line, it is common for blocks
+        # to start with a linebreak, while the previous block does not contain a
+        # linebreak
+        # in such a case we want to move the linebreak into the previous block
+        # formatters reset attributes at each line's end
+        if (lines and lines[0] == '\n' and
+            ret and not ret[-1][1].endswith('\n')):
+            ret[-1] = (ret[-1][0], ret[-1][1] + lines[0])
+            lines   = lines[1:]
+
+        for line in lines:
+            ret.append((state.to_urwid(preserve_bg), line))
+
+    start = 0
+    for m in re.finditer(esc_pattern, text):
+        # add text up to current esc-code using current terminal state
+        append_themed_block(text[start:m.start()])
+        start = m.end()
+
+        # update terminal state with this escape code
+        state.process_sgr_escape(*m.groups())
+
+    append_themed_block(text[start:])
+    return ret