Use bytes for sys.stdout.write in PY2. Closes pre-commit#161.

asottile · asottile · commit ecf82ed5e00e · 2014-09-02T16:11:45.000-07:00
diff --git a/pre_commit/commands/run.py b/pre_commit/commands/run.py
@@ -9,6 +9,7 @@
 from pre_commit import color
 from pre_commit.logging_handler import LoggingHandler
 from pre_commit.output import get_hook_message
+from pre_commit.output import sys_stdout_write_wrapper
 from pre_commit.staged_files_only import staged_files_only
 from pre_commit.util import noop_context
 
@@ -125,7 +126,7 @@ def _has_unmerged_paths(runner):
     return bool(stdout.strip())
 
 
-def run(runner, args, write=sys.stdout.write, environ=os.environ):
+def run(runner, args, write=sys_stdout_write_wrapper, environ=os.environ):
     # Set up our logging handler
     logger.addHandler(LoggingHandler(args.color, write=write))
     logger.setLevel(logging.INFO)
diff --git a/pre_commit/output.py b/pre_commit/output.py
@@ -1,8 +1,10 @@
 from __future__ import unicode_literals
 
 import subprocess
+import sys
 
 from pre_commit import color
+from pre_commit import five
 
 
 # TODO: smell: import side-effects
@@ -70,3 +72,14 @@ def get_hook_message(
             postfix,
             color.format_color(end_msg, end_color, use_color),
         )
+
+
+def sys_stdout_write_wrapper(s, stream=sys.stdout):
+    """Python 2.6 chokes on unicode being passed to sys.stdout.write.
+
+    This is an adapter because PY2 is ok with bytes and PY3 requires text.
+    """
+    assert type(s) is five.text
+    if five.PY2:  # pragma: no cover (PY2)
+        s = s.encode('UTF-8')
+    stream.write(s)
diff --git a/tests/commands/run_test.py b/tests/commands/run_test.py
@@ -1,12 +1,15 @@
+# -*- coding: UTF-8 -*-
 from __future__ import unicode_literals
 
 import io
 import mock
 import os
 import os.path
 import pytest
+import subprocess
 from plumbum import local
 
+from pre_commit.commands.install_uninstall import install
 from pre_commit.commands.run import _get_skips
 from pre_commit.commands.run import _has_unmerged_paths
 from pre_commit.commands.run import run
@@ -225,3 +228,30 @@ def test_multiple_hooks_same_id(
     ret, output = _do_run(repo_with_passing_hook, _get_opts())
     assert ret == 0
     assert output.count('Bash hook') == 2
+
+
+def test_stdout_write_bug_py26(
+        repo_with_failing_hook, mock_out_store_directory, tmpdir_factory,
+):
+    with local.cwd(repo_with_failing_hook):
+        # Add bash hook on there again
+        with io.open('.pre-commit-config.yaml', 'a+') as config_file:
+            config_file.write('        args: ["☃"]\n')
+        local['git']('add', '.pre-commit-config.yaml')
+        stage_a_file()
+
+        install(Runner(repo_with_failing_hook))
+
+        # Don't want to write to home directory
+        env = dict(os.environ, **{'PRE_COMMIT_HOME': tmpdir_factory.get()})
+        # Have to use subprocess because pytest monkeypatches sys.stdout
+        _, stdout, _ = local['git'].run(
+            ('commit', '-m', 'Commit!'),
+            # git commit puts pre-commit to stderr
+            stderr=subprocess.STDOUT,
+            env=env,
+            retcode=None,
+        )
+        assert 'UnicodeEncodeError' not in stdout
+        # Doesn't actually happen, but a reasonable assertion
+        assert 'UnicodeDecodeError' not in stdout
diff --git a/tests/output_test.py b/tests/output_test.py
@@ -1,9 +1,11 @@
 from __future__ import unicode_literals
 
+import mock
 import pytest
 
 from pre_commit import color
 from pre_commit.output import get_hook_message
+from pre_commit.output import sys_stdout_write_wrapper
 
 
 @pytest.mark.parametrize(
@@ -77,3 +79,9 @@ def test_make_sure_postfix_is_not_colored():
     assert ret == (
         'start' + '.' * 6 + 'post ' + color.RED + 'end' + color.NORMAL + '\n'
     )
+
+
+def test_sys_stdout_write_wrapper_writes():
+    fake_stream = mock.Mock()
+    sys_stdout_write_wrapper('hello world', fake_stream)
+    assert fake_stream.write.call_count == 1
