Reland "[llvm-lit] Process ANSI color codes in test output when forma… (#108107)

…tting" (#108104)"

This recommits 0f56ba1 (reverted by
6007ad7). In the original patch
llvm/utils/lit/tests/escape-color.py failed on Windows because it diffed
llvm-lit output with a file containing '\n' newlines rather than '\r\n'.
This issue is avoided by calling 'diff --strip-trailing-cr'.

Original description below:
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

Author: hnrklssn

llvm/utils/lit/lit/TestRunner.py

@@ -1017,6 +1017,20 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
     return exitCode
 
 
+def findColor(line, curr_color):
+    start = line.rfind("\33[")
+    if start == -1:
+        return curr_color
+    end = line.find("m", start + 2)
+    if end == -1:
+        return curr_color
+    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
+
+
 def formatOutput(title, data, limit=None):
     if not data.strip():
         return ""
@@ -1027,8 +1041,18 @@ def formatOutput(title, data, limit=None):
         msg = ""
     ndashes = 30
     # fmt: off
-    out =  f"# .---{title}{'-' * (ndashes - 4 - len(title))}\n"
-    out += f"# | " + "\n# | ".join(data.splitlines()) + "\n"
+    out = f"# .---{title}{'-' * (ndashes - 4 - len(title))}\n"
+    curr_color = None
+    for line in data.splitlines():
+        if curr_color:
+            out += "\33[0m"
+        out += "# | "
+        if curr_color:
+            out += curr_color
+        out += line + "\n"
+        curr_color = findColor(line, curr_color)
+    if curr_color:
+        out += "\33[0m"  # prevent unterminated formatting from leaking
     out += f"# `---{msg}{'-' * (ndashes - 4 - len(msg))}\n"
     # fmt: on
     return out

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)
+# `-----------------------------
+
+--

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)

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
+

llvm/utils/lit/tests/escape-color.py

@@ -0,0 +1,4 @@
+# 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 --strip-trailing-cr %{inputs}/escape-color/color-escaped.txt %t