Handle unicode output properly in CLI

2 files changed, 19 insertions, 7 deletions

diff --git a/scripts/mammoth b/scripts/mammoth
index f4470f7..16563e9 100644
--- a/scripts/mammoth
+++ b/scripts/mammoth
@@ -22,19 +22,31 @@ def main():
             for message in result.messages:
                 sys.stderr.write(message.message)
                 sys.stderr.write("\n")
-            output.write(result.value.encode("utf8"))
+            output.write(result.value)
 
 
 @contextlib.contextmanager
 def _open_output(name):
     if name is None:
         try:
-            yield sys.stdout
+            yield Output(sys.stdout)
         finally:
             sys.stdout.flush()
     else:
         with open(name, "w") as output:
-            yield output
+            yield Output(output)
+
+
+class Output(object):
+    def __init__(self, fileobj):
+        self._fileobj = fileobj
+    
+    def write(self, value):
+        if sys.version_info[0] <= 2: 
+            self._fileobj.write(value.encode("utf8"))
+        else:
+            self._fileobj.write(value)
+            
 
 
 def _parse_args():
diff --git a/tests/cli_tests.py b/tests/cli_tests.py
index 4ce1d87..a54349b 100644
--- a/tests/cli_tests.py
+++ b/tests/cli_tests.py
@@ -13,8 +13,8 @@ _local = spur.LocalShell()
 def html_is_printed_to_stdout_if_output_file_is_not_set():
     docx_path = test_path("single-paragraph.docx")
     result = _local.run(["mammoth", docx_path])
-    assert_equal("", result.stderr_output)
-    assert_equal("<p>Walking on imported air</p>", result.output)
+    assert_equal(b"", result.stderr_output)
+    assert_equal(b"<p>Walking on imported air</p>", result.output)
 
 
 @istest
@@ -23,8 +23,8 @@ def html_is_written_to_file_if_output_file_is_set():
         output_path = os.path.join(temp_dir.path, "output.html")
         docx_path = test_path("single-paragraph.docx")
         result = _local.run(["mammoth", docx_path, output_path])
-        assert_equal("", result.stderr_output)
-        assert_equal("", result.output)
+        assert_equal(b"", result.stderr_output)
+        assert_equal(b"", result.output)
         with open(output_path) as output_file:
             assert_equal("<p>Walking on imported air</p>", output_file.read())