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

llvm/utils/lit/lit/TestRunner.py

@@ -1005,6 +1005,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 ""
@@ -1015,8 +1029,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 %{inputs}/escape-color/color-escaped.txt %t