[llvm-lit] Process ANSI color codes in test output when formatting

[hnrklssn]

Test output that carried color across newlines previously resulted in the formatting around the output also being colored. Detect the current ANSI color and reset it when printing formatting, and then reapply it. As an added bonus an unterminated color code is also detected, preventing it from leaking out into the rest of the terminal.

Fixes #106633

[llvmbot]

@llvm/pr-subscribers-testing-tools

Author: Henrik G. Olsson (hnrklssn)

Changes

Test output that carried color across newlines previously resulted in the formatting around the output also being colored. Detect the current ANSI color and reset it when printing formatting, and then reapply it. As an added bonus an unterminated color code is also detected, preventing it from leaking out into the rest of the terminal.

Fixes #106633

---

Full diff: https://github.com/llvm/llvm-project/pull/106776.diff

5 Files Affected:

• (modified) llvm/utils/lit/lit/TestRunner.py (+32-1)
• (added) llvm/utils/lit/tests/Inputs/escape-color/color-escaped.txt (+10)
• (added) llvm/utils/lit/tests/Inputs/escape-color/color.txt (+6)
• (added) llvm/utils/lit/tests/Inputs/escape-color/lit.cfg (+8)
• (added) llvm/utils/lit/tests/escape-color.py (+5)

diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index da7fa86fd39173..329603617a0ab2 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1004,6 +1004,37 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
 
     return exitCode
 
+def reFindLast(r, s):
+    matches = r.findall(s)
+    if not matches:
+        return None
+    return matches[-1]
+
+def findColor(s, curr_color):
+    m = reFindLast(color_re, s)
+    if not m:
+        return curr_color
+    # "\33[0m" means "reset all formatting". Sometimes the 0 is skipped.
+    if m == "\33[m" or m == "\33[0m":
+        return None
+    return m
+
+def escapeColor(s, f):
+    if not f:
+        return s
+    return f"\33[0m{s}{f}"
+
+color_re = re.compile("\33\\[[^m]*m")
+def formatLines(lines):
+    curr_color = None
+    out = ""
+    for line in lines:
+        out += escapeColor("# | ", curr_color)
+        out += line + "\n"
+        curr_color = findColor(line, curr_color)
+    if curr_color:
+        out += "\33[0m" # prevent unterminated formatting from leaking
+    return out
 
 def formatOutput(title, data, limit=None):
     if not data.strip():
@@ -1016,7 +1047,7 @@ def formatOutput(title, data, limit=None):
     ndashes = 30
     # fmt: off
     out =  f"# .---{title}{'-' * (ndashes - 4 - len(title))}\n"
-    out += f"# | " + "\n# | ".join(data.splitlines()) + "\n"
+    out += formatLines(data.splitlines())
     out += f"# `---{msg}{'-' * (ndashes - 4 - len(msg))}\n"
     # fmt: on
     return out
diff --git a/llvm/utils/lit/tests/Inputs/escape-color/color-escaped.txt b/llvm/utils/lit/tests/Inputs/escape-color/color-escaped.txt
new file mode 100644
index 00000000000000..e7a33e380b351c
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/escape-color/color-escaped.txt
@@ -0,0 +1,10 @@
+# .---command stdout------------
+# | # RUN: cat %s
+# | �[31mred
+�[0m# | �[31mstill red�(B�[m
+# | plain
+# | �[32mgreen
+�[0m# | �[32mstill green (never terminated)
+�[0m# `-----------------------------
+
+--
diff --git a/llvm/utils/lit/tests/Inputs/escape-color/color.txt b/llvm/utils/lit/tests/Inputs/escape-color/color.txt
new file mode 100644
index 00000000000000..15ffc22d134f0f
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/escape-color/color.txt
@@ -0,0 +1,6 @@
+# RUN: cat %s
+�[31mred
+still red�(B�[m
+plain
+�[32mgreen
+still green (never terminated)
diff --git a/llvm/utils/lit/tests/Inputs/escape-color/lit.cfg b/llvm/utils/lit/tests/Inputs/escape-color/lit.cfg
new file mode 100644
index 00000000000000..36f4eb69d4858e
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/escape-color/lit.cfg
@@ -0,0 +1,8 @@
+import lit.formats
+
+config.name = "escape-color"
+config.suffixes = [".txt"]
+config.test_format = lit.formats.ShTest()
+config.test_source_root = None
+config.test_exec_root = None
+
diff --git a/llvm/utils/lit/tests/escape-color.py b/llvm/utils/lit/tests/escape-color.py
new file mode 100644
index 00000000000000..35b34e47e45586
--- /dev/null
+++ b/llvm/utils/lit/tests/escape-color.py
@@ -0,0 +1,5 @@
+# cut off the first 9 lines to avoid absolute file paths in the output
+# then keep only the next 10 lines to avoid test timing in the output
+# RUN: %{lit} %{inputs}/escape-color/color.txt -a | tail -n +10 | head -n 10 > %t
+# RUN: diff %{inputs}/escape-color/color-escaped.txt %t
+

[github-actions]

⚠️ Python code formatter, darker found issues in your code. ⚠️

You can test this locally with the following command:

darker --check --diff -r 9e1f1cfa59c4513798de5c326d77e1e5912b1634...a9117b4d2927544ab870e08d8e1d0f4d6d626601 llvm/utils/lit/tests/escape-color.py llvm/utils/lit/lit/TestRunner.py

View the diff from darker here.

--- lit/TestRunner.py	2024-09-07 01:23:47.000000 +0000
+++ lit/TestRunner.py	2024-09-07 01:33:22.295326 +0000
@@ -1007,14 +1007,14 @@
 
 def findColor(line, curr_color):
     start = line.rfind("\33[")
     if start == -1:
         return curr_color
-    end = line.find("m", start+2)
+    end = line.find("m", start + 2)
     if end == -1:
         return curr_color
-    match = line[start:end+1]
+    match = line[start : end + 1]
     # "\33[0m" means "reset all formatting". Sometimes the 0 is skipped.
     if match == "\33[m" or match == "\33[0m":
         return None
     return match
 

[arichardson]

Overall this makes sense to me just a few stylistic comments.

[hnrklssn]

Overall this makes sense to me just a few stylistic comments.

Cool, I've inlined most of the functions and switched to find + slice based parsing instead of regex.

[dyung]

Hi @hnrklssn, the test you added, escape-color.py seems to be failing on a Windows bot:

https://lab.llvm.org/buildbot/#/builders/46/builds/4703

Can you take a look and revert if you need time to investigate?

[hnrklssn]

Hi @hnrklssn, the test you added, escape-color.py seems to be failing on a Windows bot:

https://lab.llvm.org/buildbot/#/builders/46/builds/4703

Can you take a look and revert if you need time to investigate?

Yeah I saw that but I can't figure out why it's failing, so I'll revert: the two inputs to diff look the same to me...
I'll see if I can test it out on a Windows machine later and see what's going on there.

[hnrklssn]

Oh wait it's probably because of \r\n vs \n. I'll try again with diff --strip-trailing-cr.