From 59b98a7f33115be45dab595cc93837d23ae523e2 Mon Sep 17 00:00:00 2001
From: Ivan Habunek <ivan@habunek.com>
Date: Fri, 2 Dec 2022 07:40:47 +0100
Subject: [PATCH] Improve colorize

---
 tests/test_output.py |  26 +++++++++++
 toot/output.py       | 108 ++++++++++++++++++++++++++++++++-----------
 2 files changed, 107 insertions(+), 27 deletions(-)
 create mode 100644 tests/test_output.py

diff --git a/tests/test_output.py b/tests/test_output.py
new file mode 100644
index 0000000..cc31e5c
--- /dev/null
+++ b/tests/test_output.py
@@ -0,0 +1,26 @@
+from toot.output import colorize, strip_tags, STYLES
+
+reset = STYLES["reset"]
+red = STYLES["red"]
+green = STYLES["green"]
+bold = STYLES["bold"]
+
+
+def test_colorize():
+    assert colorize("foo") == "foo"
+    assert colorize("<red>foo</red>") == f"{red}foo{reset}{reset}"
+    assert colorize("foo <red>bar</red> baz") == f"foo {red}bar{reset} baz{reset}"
+    assert colorize("foo <red bold>bar</red bold> baz") == f"foo {red}{bold}bar{reset} baz{reset}"
+    assert colorize("foo <red bold>bar</red> baz") == f"foo {red}{bold}bar{reset}{bold} baz{reset}"
+    assert colorize("foo <red bold>bar</> baz") == f"foo {red}{bold}bar{reset} baz{reset}"
+    assert colorize("<red>foo<bold>bar</bold>baz</red>") == f"{red}foo{bold}bar{reset}{red}baz{reset}{reset}"
+
+
+def test_strip_tags():
+    assert strip_tags("foo") == "foo"
+    assert strip_tags("<red>foo</red>") == "foo"
+    assert strip_tags("foo <red>bar</red> baz") == "foo bar baz"
+    assert strip_tags("foo <red bold>bar</red bold> baz") == "foo bar baz"
+    assert strip_tags("foo <red bold>bar</red> baz") == "foo bar baz"
+    assert strip_tags("foo <red bold>bar</> baz") == "foo bar baz"
+    assert strip_tags("<red>foo<bold>bar</bold>baz</red>") == "foobarbaz"
diff --git a/toot/output.py b/toot/output.py
index 20ef3bb..0c74f61 100644
--- a/toot/output.py
+++ b/toot/output.py
@@ -13,39 +13,93 @@ from toot.utils import get_text, parse_html
 from toot.wcstring import wc_wrap
 
 
-START_CODES = {
-    'red': '\033[31m',
-    'green': '\033[32m',
-    'yellow': '\033[33m',
-    'blue': '\033[34m',
-    'magenta': '\033[35m',
-    'cyan': '\033[36m',
+STYLES = {
+    'reset': '\033[0m',
+    'bold': '\033[1m',
+    'dim': '\033[2m',
+    'italic': '\033[3m',
+    'underline': '\033[4m',
+    'red': '\033[91m',
+    'green': '\033[92m',
+    'yellow': '\033[93m',
+    'blue': '\033[94m',
+    'magenta': '\033[95m',
+    'cyan': '\033[96m',
 }
 
-END_CODE = '\033[0m'
-
-START_PATTERN = "<(" + "|".join(START_CODES.keys()) + ")>"
-
-END_PATTERN = "</(" + "|".join(START_CODES.keys()) + ")>"
+STYLE_TAG_PATTERN = re.compile(r"""
+    (?<!\\)     # not preceeded by a backslash - allows escaping
+    <           # literal
+    (/)?        # optional closing - first group
+    (.*?)       # style names - ungreedy - second group
+    >           # literal
+""", re.X)
 
 
-def start_code(match):
-    name = match.group(1)
-    return START_CODES[name]
+def colorize(message):
+    """
+    Replaces style tags in `message` with ANSI escape codes.
+
+    Markup is inspired by HTML, but you can use multiple words pre tag, e.g.:
+
+        <red bold>alert!</red bold> a thing happened
+
+    Empty closing tag will reset all styes:
+
+        <red bold>alert!</> a thing happened
+
+    Styles can be nested:
+
+        <red>red <underline>red and underline</underline> red</red>
+    """
+
+    def _codes(styles):
+        for style in styles:
+            yield STYLES.get(style, "")
+
+    def _generator(message):
+        # A list is used instead of a set because we want to keep style order
+        # This allows nesting colors, e.g. "<blue>foo<red>bar</red>baz</blue>"
+        position = 0
+        active_styles = []
+
+        for match in re.finditer(STYLE_TAG_PATTERN, message):
+            is_closing = bool(match.group(1))
+            styles = match.group(2).strip().split()
+
+            start, end = match.span()
+            # Replace backslash for escaped <
+            yield message[position:start].replace("\\<", "<")
+
+            if is_closing:
+                yield STYLES["reset"]
+
+                # Empty closing tag resets all styles
+                if styles == []:
+                    active_styles = []
+                else:
+                    active_styles = [s for s in active_styles if s not in styles]
+                    yield from _codes(active_styles)
+            else:
+                active_styles = active_styles + styles
+                yield from _codes(styles)
+
+            position = end
+
+        if position == 0:
+            # Nothing matched, yield the original string
+            yield message
+        else:
+            # Yield the remaining fragment
+            yield message[position:]
+            # Reset styles at the end to prevent leaking
+            yield STYLES["reset"]
+
+    return "".join(_generator(message))
 
 
-def colorize(text):
-    text = re.sub(START_PATTERN, start_code, text)
-    text = re.sub(END_PATTERN, END_CODE, text)
-
-    return text
-
-
-def strip_tags(text):
-    text = re.sub(START_PATTERN, '', text)
-    text = re.sub(END_PATTERN, '', text)
-
-    return text
+def strip_tags(message):
+    return re.sub(STYLE_TAG_PATTERN, "", message)
 
 
 def use_ansi_color():