[llvm] [llvm-lit] Process ANSI color codes in test output when formatting (PR #106776)

[Henrik G. Olsson via llvm-commits]

https://github.com/hnrklssn created https://github.com/llvm/llvm-project/pull/106776

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

>From e8733617420b4d67133cbd23312afa0d5639220e Mon Sep 17 00:00:00 2001
From: "Henrik G. Olsson" <h_olsson at apple.com>
Date: Fri, 30 Aug 2024 01:13:51 -0700
Subject: [PATCH] [llvm-lit] Process ANSI color codes in test output when
 formatting

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
---
 llvm/utils/lit/lit/TestRunner.py              | 33 ++++++++++++++++++-
 .../Inputs/escape-color/color-escaped.txt     | 10 ++++++
 .../lit/tests/Inputs/escape-color/color.txt   |  6 ++++
 .../lit/tests/Inputs/escape-color/lit.cfg     |  8 +++++
 llvm/utils/lit/tests/escape-color.py          |  5 +++
 5 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 llvm/utils/lit/tests/Inputs/escape-color/color-escaped.txt
 create mode 100644 llvm/utils/lit/tests/Inputs/escape-color/color.txt
 create mode 100644 llvm/utils/lit/tests/Inputs/escape-color/lit.cfg
 create mode 100644 llvm/utils/lit/tests/escape-color.py

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
+# | red
+# | still red(B
+# | plain
+# | green
+# | still green (never terminated)
+# `-----------------------------
+
+--
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
+red
+still red(B
+plain
+green
+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
+