[lit] Don't fail when printing test output with special chars

This addresses a UnicodeEncodeError when using Python 3.6.5 in Windows
10.

Reviewed By: rnk

Differential Revision: https://reviews.llvm.org/D69207

4 files changed, 35 insertions, 3 deletions

diff --git a/llvm/utils/lit/lit/display.py b/llvm/utils/lit/lit/display.py
index 5c6d6b6e08e..26700134d2e 100644
--- a/llvm/utils/lit/lit/display.py
+++ b/llvm/utils/lit/lit/display.py
@@ -70,7 +70,20 @@ class ProgressDisplay(object):
             if test.result.code.isFailure:
                 print("%s TEST '%s' FAILED %s" % ('*'*20, test.getFullName(),
                                                   '*'*20))
-            print(test.result.output)
+            out = test.result.output
+            # Encode/decode so that, when using Python 3.6.5 in Windows 10,
+            # print(out) doesn't raise UnicodeEncodeError if out contains
+            # special characters.  However, Python 2 might try to decode
+            # as part of the encode call if out is already encoded, so skip
+            # encoding if it raises UnicodeDecodeError.
+            if sys.stdout.encoding:
+                try:
+                    out = out.encode(encoding=sys.stdout.encoding,
+                                     errors="replace")
+                except UnicodeDecodeError:
+                    pass
+                out = out.decode(encoding=sys.stdout.encoding)
+            print(out)
             print("*" * 20)
 
         # Report test metrics, if present.
diff --git a/llvm/utils/lit/tests/Inputs/shtest-shell/stdout-encoding.txt b/llvm/utils/lit/tests/Inputs/shtest-shell/stdout-encoding.txt
new file mode 100644
index 00000000000..8a891c8c624
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/shtest-shell/stdout-encoding.txt
@@ -0,0 +1,7 @@
+# Check that lit doesn't fail when printing special characters in its test
+# results.
+
+# RUN: cat diff-in.bin
+
+# Fail so lit will print output.
+# RUN: false
diff --git a/llvm/utils/lit/tests/max-failures.py b/llvm/utils/lit/tests/max-failures.py
index 6ad37533692..f661980ac2b 100644
--- a/llvm/utils/lit/tests/max-failures.py
+++ b/llvm/utils/lit/tests/max-failures.py
@@ -8,7 +8,7 @@
 #
 # END.
 
-# CHECK: Failing Tests (32)
+# CHECK: Failing Tests (33)
 # CHECK: Failing Tests (1)
 # CHECK: Failing Tests (2)
 # CHECK: error: argument --max-failures: requires positive integer, but found '0'
diff --git a/llvm/utils/lit/tests/shtest-shell.py b/llvm/utils/lit/tests/shtest-shell.py
index 1fc8b30c7ee..cef3728dd00 100644
--- a/llvm/utils/lit/tests/shtest-shell.py
+++ b/llvm/utils/lit/tests/shtest-shell.py
@@ -413,5 +413,17 @@
 # CHECK: PASS: shtest-shell :: rm-unicode-0.txt
 # CHECK: PASS: shtest-shell :: sequencing-0.txt
 # CHECK: XFAIL: shtest-shell :: sequencing-1.txt
+
+# CHECK: FAIL: shtest-shell :: stdout-encoding.txt
+# CHECK: *** TEST 'shtest-shell :: stdout-encoding.txt' FAILED ***
+# CHECK: $ "cat" "diff-in.bin"
+# CHECK: # command output:
+# CHECK-NEXT: {{^.f.o.o.$}}
+# CHECK-NEXT: {{^.b.a.r..}}
+# CHECK-NEXT: {{^.b.a.z.$}}
+# CHECK-NOT: error
+# CHECK: $ "false"
+# CHECK: ***
+
 # CHECK: PASS: shtest-shell :: valid-shell.txt
-# CHECK: Failing Tests (32)
+# CHECK: Failing Tests (33)