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

[Henrik G. Olsson via llvm-commits]

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

…tting" (#108104)"

This recommits 0f56ba13bff7ab72bfafcf7c5cf9e5b8bd16d895 (reverted by 6007ad79afeffb1288781b4a7241290386293aff). 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

>From 6a04b1f26ce36e183a44a7b822f75980e76a3521 Mon Sep 17 00:00:00 2001
From: "Henrik G. Olsson" <h_olsson at apple.com>
Date: Tue, 10 Sep 2024 15:31:13 -0700
Subject: [PATCH] Reland "[llvm-lit] Process ANSI color codes in test output
 when formatting" (#108104)"

This recommits 0f56ba13bff7ab72bfafcf7c5cf9e5b8bd16d895 (reverted by
6007ad79afeffb1288781b4a7241290386293aff). 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
---
 llvm/utils/lit/lit/TestRunner.py              | 28 +++++++++++++++++--
 .../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          |  4 +++
 5 files changed, 54 insertions(+), 2 deletions(-)
 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 19f35fc7e212f3..a2c76d41a43e07 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/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
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..8fdda3553da399
--- /dev/null
+++ b/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