diff options
| -rw-r--r-- | scripts/mammoth | 18 | ||||
| -rw-r--r-- | tests/cli_tests.py | 8 |
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()) |