testsuite: Factor out terminal coloring

4 files changed, 36 insertions, 19 deletions

diff --git a/testsuite/driver/runtests.py b/testsuite/driver/runtests.py
index 5ef6ca27bf..9b26e1bf75 100644
--- a/testsuite/driver/runtests.py
+++ b/testsuite/driver/runtests.py
@@ -27,6 +27,8 @@ from testutil import getStdout, Watcher, str_warn, str_info
 from testglobals import getConfig, ghc_env, getTestRun, TestConfig, TestOptions, brokens
 from perf_notes import MetricChange, inside_git_repo, is_worktree_dirty, format_perf_stat
 from junit import junit
+import term_color
+from term_color import Color, colored
 import cpu_features
 
 # Readline sometimes spews out ANSI escapes for some values of TERM,
@@ -213,6 +215,7 @@ def supports_colors():
     return True
 
 config.supports_colors = supports_colors()
+term_color.enable_color = config.supports_colors
 
 # This has to come after arg parsing as the args can change the compiler
 get_compiler_info()
diff --git a/testsuite/driver/term_color.py b/testsuite/driver/term_color.py
new file mode 100644
index 0000000000..ba9ff7941a
--- /dev/null
+++ b/testsuite/driver/term_color.py
@@ -0,0 +1,20 @@
+from enum import Enum
+
+enable_color = True
+
+class Color(Enum):
+    BLACK   = 30
+    RED     = 31
+    GREEN   = 32
+    YELLOW  = 33
+    BLUE    = 34
+    MAGENTA = 35
+    CYAN    = 36
+    WHITE   = 37
+
+def colored(color: Color, s: str) -> str:
+    if enable_color:
+        return '\033[1m\033[{}m{}\033[0m'.format(color.value, s)
+    else:
+        return s
+
diff --git a/testsuite/driver/testlib.py b/testsuite/driver/testlib.py
index 07206799c1..9d24951361 100644
--- a/testsuite/driver/testlib.py
+++ b/testsuite/driver/testlib.py
@@ -21,8 +21,9 @@ import subprocess
 from testglobals import config, ghc_env, default_testopts, brokens, t, \
                         TestRun, TestResult, TestOptions
 from testutil import strip_quotes, lndir, link_or_copy_file, passed, \
-                     failBecause, str_fail, str_pass, testing_metrics, \
+                     failBecause, testing_metrics, \
                      PassFail
+from term_color import Color, colored
 import testutil
 from cpu_features import have_cpu_feature
 import perf_notes as Perf
@@ -2430,18 +2431,16 @@ def summary(t: TestRun, file: TextIO, short=False, color=False) -> None:
         # Only print the list of unexpected tests above.
         return
 
-    colorize = lambda s: s
-    if color:
-        if len(t.unexpected_failures) > 0 or \
-            len(t.unexpected_stat_failures) > 0 or \
-            len(t.unexpected_passes) > 0 or \
-            len(t.framework_failures) > 0:
-            colorize = str_fail
-        else:
-            colorize = str_pass
+    if len(t.unexpected_failures) > 0 or \
+        len(t.unexpected_stat_failures) > 0 or \
+        len(t.unexpected_passes) > 0 or \
+        len(t.framework_failures) > 0:
+        summary_color = Color.RED
+    else:
+        summary_color = Color.GREEN
 
     assert t.start_time is not None
-    file.write(colorize('SUMMARY') + ' for test run started at '
+    file.write(colored(summary_color, 'SUMMARY') + ' for test run started at '
                + t.start_time.strftime("%c %Z") + '\n'
                + str(datetime.datetime.now() - t.start_time).rjust(8)
                + ' spent to go through\n'
diff --git a/testsuite/driver/testutil.py b/testsuite/driver/testutil.py
index 071b641b7f..19cdbf3625 100644
--- a/testsuite/driver/testutil.py
+++ b/testsuite/driver/testutil.py
@@ -3,6 +3,7 @@ import platform
 import subprocess
 import shutil
 from pathlib import Path, PurePath
+from term_color import Color, colored
 
 import threading
 
@@ -37,17 +38,11 @@ def strip_quotes(s: str) -> str:
     # Don't wrap commands to subprocess.call/Popen in quotes.
     return s.strip('\'"')
 
-def str_fail(s: str) -> str:
-    return '\033[1m\033[31m' + s + '\033[0m'
-
-def str_pass(s: str) -> str:
-    return '\033[1m\033[32m' + s + '\033[0m'
-
 def str_warn(s: str) -> str:
-    return '\033[1m\033[33m' + s + '\033[0m'
+    return colored(Color.YELLOW, s)
 
 def str_info(s: str) -> str:
-    return '\033[1m\033[34m' + s + '\033[0m'
+    return colored(Color.BLUE, s)
 
 def getStdout(cmd_and_args: List[str]):
     # Can't use subprocess.check_output, since we also verify that