show color in hook outputs when attached to a tty

asottile · asottile · commit 7c3404ef1f75 · 2019-10-12T23:17:36.000-07:00
diff --git a/pre_commit/commands/run.py b/pre_commit/commands/run.py
@@ -73,7 +73,7 @@ def _hook_msg_start(hook, verbose):
 NO_FILES = '(no files to check)'
 
 
-def _run_single_hook(classifier, hook, args, skips, cols):
+def _run_single_hook(classifier, hook, args, skips, cols, use_color):
     filenames = classifier.filenames_for_hook(hook)
 
     if hook.language == 'pcre':
@@ -118,7 +118,8 @@ def _run_single_hook(classifier, hook, args, skips, cols):
     sys.stdout.flush()
 
     diff_before = cmd_output_b('git', 'diff', '--no-ext-diff', retcode=None)
-    retcode, out = hook.run(tuple(filenames) if hook.pass_filenames else ())
+    filenames = tuple(filenames) if hook.pass_filenames else ()
+    retcode, out = hook.run(filenames, use_color)
     diff_after = cmd_output_b('git', 'diff', '--no-ext-diff', retcode=None)
 
     file_modifications = diff_before != diff_after
@@ -203,7 +204,9 @@ def _run_hooks(config, hooks, args, environ):
     classifier = Classifier(filenames)
     retval = 0
     for hook in hooks:
-        retval |= _run_single_hook(classifier, hook, args, skips, cols)
+        retval |= _run_single_hook(
+            classifier, hook, args, skips, cols, args.color,
+        )
         if retval and config['fail_fast']:
             break
     if retval and args.show_diff_on_failure and git.has_diff():
diff --git a/pre_commit/languages/all.py b/pre_commit/languages/all.py
@@ -38,16 +38,17 @@
 #         version - A version specified in the hook configuration or 'default'.
 #     """
 #
-# def run_hook(hook, file_args):
+# def run_hook(hook, file_args, color):
 #     """Runs a hook and returns the returncode and output of running that
 #     hook.
 #
 #     Args:
 #         hook - `Hook`
 #         file_args - The files to be run
+#         color - whether the hook should be given a pty (when supported)
 #
 #     Returns:
-#         (returncode, stdout, stderr)
+#         (returncode, output)
 #     """
 
 languages = {
diff --git a/pre_commit/languages/docker.py b/pre_commit/languages/docker.py
@@ -95,15 +95,15 @@ def docker_cmd():  # pragma: windows no cover
     )
 
 
-def run_hook(hook, file_args):  # pragma: windows no cover
+def run_hook(hook, file_args, color):  # pragma: windows no cover
     assert_docker_available()
     # Rebuild the docker image in case it has gone missing, as many people do
     # automated cleanup of docker images.
     build_docker_image(hook.prefix, pull=False)
 
-    hook_cmd = helpers.to_cmd(hook)
-    entry_exe, cmd_rest = hook_cmd[0], hook_cmd[1:]
+    hook_cmd = hook.cmd
+    entry_exe, cmd_rest = hook.cmd[0], hook_cmd[1:]
 
     entry_tag = ('--entrypoint', entry_exe, docker_tag(hook.prefix))
     cmd = docker_cmd() + entry_tag + cmd_rest
-    return helpers.run_xargs(hook, cmd, file_args)
+    return helpers.run_xargs(hook, cmd, file_args, color=color)
diff --git a/pre_commit/languages/docker_image.py b/pre_commit/languages/docker_image.py
@@ -12,7 +12,7 @@
 install_environment = helpers.no_install
 
 
-def run_hook(hook, file_args):  # pragma: windows no cover
+def run_hook(hook, file_args, color):  # pragma: windows no cover
     assert_docker_available()
-    cmd = docker_cmd() + helpers.to_cmd(hook)
-    return helpers.run_xargs(hook, cmd, file_args)
+    cmd = docker_cmd() + hook.cmd
+    return helpers.run_xargs(hook, cmd, file_args, color=color)
diff --git a/pre_commit/languages/fail.py b/pre_commit/languages/fail.py
@@ -9,7 +9,7 @@
 install_environment = helpers.no_install
 
 
-def run_hook(hook, file_args):
+def run_hook(hook, file_args, color):
     out = hook.entry.encode('UTF-8') + b'\n\n'
     out += b'\n'.join(f.encode('UTF-8') for f in file_args) + b'\n'
     return 1, out
diff --git a/pre_commit/languages/golang.py b/pre_commit/languages/golang.py
@@ -81,6 +81,6 @@ def install_environment(prefix, version, additional_dependencies):
             rmtree(pkgdir)
 
 
-def run_hook(hook, file_args):
+def run_hook(hook, file_args, color):
     with in_env(hook.prefix):
-        return helpers.run_xargs(hook, helpers.to_cmd(hook), file_args)
+        return helpers.run_xargs(hook, hook.cmd, file_args, color=color)
diff --git a/pre_commit/languages/helpers.py b/pre_commit/languages/helpers.py
@@ -3,7 +3,6 @@
 import multiprocessing
 import os
 import random
-import shlex
 
 import six
 
@@ -25,10 +24,6 @@ def environment_dir(ENVIRONMENT_DIR, language_version):
         return '{}-{}'.format(ENVIRONMENT_DIR, language_version)
 
 
-def to_cmd(hook):
-    return tuple(shlex.split(hook.entry)) + tuple(hook.args)
-
-
 def assert_version_default(binary, version):
     if version != C.DEFAULT:
         raise AssertionError(
@@ -83,8 +78,9 @@ def _shuffled(seq):
     return seq
 
 
-def run_xargs(hook, cmd, file_args):
+def run_xargs(hook, cmd, file_args, **kwargs):
     # Shuffle the files so that they more evenly fill out the xargs partitions,
     # but do it deterministically in case a hook cares about ordering.
     file_args = _shuffled(file_args)
-    return xargs(cmd, file_args, target_concurrency=target_concurrency(hook))
+    kwargs['target_concurrency'] = target_concurrency(hook)
+    return xargs(cmd, file_args, **kwargs)
diff --git a/pre_commit/languages/node.py b/pre_commit/languages/node.py
@@ -78,6 +78,6 @@ def install_environment(
             )
 
 
-def run_hook(hook, file_args):  # pragma: windows no cover
+def run_hook(hook, file_args, color):  # pragma: windows no cover
     with in_env(hook.prefix, hook.language_version):
-        return helpers.run_xargs(hook, helpers.to_cmd(hook), file_args)
+        return helpers.run_xargs(hook, hook.cmd, file_args, color=color)
diff --git a/pre_commit/languages/pcre.py b/pre_commit/languages/pcre.py
@@ -13,10 +13,10 @@
 install_environment = helpers.no_install
 
 
-def run_hook(hook, file_args):
+def run_hook(hook, file_args, color):
     # For PCRE the entry is the regular expression to match
     cmd = (GREP, '-H', '-n', '-P') + tuple(hook.args) + (hook.entry,)
 
     # Grep usually returns 0 for matches, and nonzero for non-matches so we
     # negate it here.
-    return xargs(cmd, file_args, negate=True)
+    return xargs(cmd, file_args, negate=True, color=color)
diff --git a/pre_commit/languages/pygrep.py b/pre_commit/languages/pygrep.py
@@ -44,9 +44,9 @@ def _process_filename_at_once(pattern, filename):
     return retv
 
 
-def run_hook(hook, file_args):
+def run_hook(hook, file_args, color):
     exe = (sys.executable, '-m', __name__) + tuple(hook.args) + (hook.entry,)
-    return xargs(exe, file_args)
+    return xargs(exe, file_args, color=color)
 
 
 def main(argv=None):
diff --git a/pre_commit/languages/python.py b/pre_commit/languages/python.py
@@ -151,9 +151,9 @@ def healthy(prefix, language_version):
             )
         return retcode == 0
 
-    def run_hook(hook, file_args):
+    def run_hook(hook, file_args, color):
         with in_env(hook.prefix, hook.language_version):
-            return helpers.run_xargs(hook, helpers.to_cmd(hook), file_args)
+            return helpers.run_xargs(hook, hook.cmd, file_args, color=color)
 
     def install_environment(prefix, version, additional_dependencies):
         additional_dependencies = tuple(additional_dependencies)
diff --git a/pre_commit/languages/ruby.py b/pre_commit/languages/ruby.py
@@ -124,6 +124,6 @@ def install_environment(
             )
 
 
-def run_hook(hook, file_args):  # pragma: windows no cover
+def run_hook(hook, file_args, color):  # pragma: windows no cover
     with in_env(hook.prefix, hook.language_version):
-        return helpers.run_xargs(hook, helpers.to_cmd(hook), file_args)
+        return helpers.run_xargs(hook, hook.cmd, file_args, color=color)
diff --git a/pre_commit/languages/rust.py b/pre_commit/languages/rust.py
@@ -89,6 +89,6 @@ def install_environment(prefix, version, additional_dependencies):
             )
 
 
-def run_hook(hook, file_args):
+def run_hook(hook, file_args, color):
     with in_env(hook.prefix):
-        return helpers.run_xargs(hook, helpers.to_cmd(hook), file_args)
+        return helpers.run_xargs(hook, hook.cmd, file_args, color=color)
diff --git a/pre_commit/languages/script.py b/pre_commit/languages/script.py
@@ -9,7 +9,7 @@
 install_environment = helpers.no_install
 
 
-def run_hook(hook, file_args):
-    cmd = helpers.to_cmd(hook)
+def run_hook(hook, file_args, color):
+    cmd = hook.cmd
     cmd = (hook.prefix.path(cmd[0]),) + cmd[1:]
-    return helpers.run_xargs(hook, cmd, file_args)
+    return helpers.run_xargs(hook, cmd, file_args, color=color)
diff --git a/pre_commit/languages/swift.py b/pre_commit/languages/swift.py
@@ -51,6 +51,6 @@ def install_environment(
         )
 
 
-def run_hook(hook, file_args):  # pragma: windows no cover
+def run_hook(hook, file_args, color):  # pragma: windows no cover
     with in_env(hook.prefix):
-        return helpers.run_xargs(hook, helpers.to_cmd(hook), file_args)
+        return helpers.run_xargs(hook, hook.cmd, file_args, color=color)
diff --git a/pre_commit/languages/system.py b/pre_commit/languages/system.py
@@ -9,5 +9,5 @@
 install_environment = helpers.no_install
 
 
-def run_hook(hook, file_args):
-    return helpers.run_xargs(hook, helpers.to_cmd(hook), file_args)
+def run_hook(hook, file_args, color):
+    return helpers.run_xargs(hook, hook.cmd, file_args, color=color)
diff --git a/pre_commit/repository.py b/pre_commit/repository.py
@@ -5,6 +5,7 @@
 import json
 import logging
 import os
+import shlex
 
 import pre_commit.constants as C
 from pre_commit import five
@@ -54,6 +55,10 @@ def _write_state(prefix, venv, state):
 class Hook(collections.namedtuple('Hook', ('src', 'prefix') + _KEYS)):
     __slots__ = ()
 
+    @property
+    def cmd(self):
+        return tuple(shlex.split(self.entry)) + tuple(self.args)
+
     @property
     def install_key(self):
         return (
@@ -95,9 +100,9 @@ def install(self):
         # Write our state to indicate we're installed
         _write_state(self.prefix, venv, _state(self.additional_dependencies))
 
-    def run(self, file_args):
+    def run(self, file_args, color):
         lang = languages[self.language]
-        return lang.run_hook(self, file_args)
+        return lang.run_hook(self, file_args, color)
 
     @classmethod
     def create(cls, src, prefix, dct):
diff --git a/pre_commit/util.py b/pre_commit/util.py
@@ -117,29 +117,28 @@ def to_text(self):
         __str__ = to_text
 
 
-def cmd_output_b(*cmd, **kwargs):
-    retcode = kwargs.pop('retcode', 0)
-
-    popen_kwargs = {
-        'stdin': subprocess.PIPE,
-        'stdout': subprocess.PIPE,
-        'stderr': subprocess.PIPE,
-    }
-
+def _cmd_kwargs(*cmd, **kwargs):
     # py2/py3 on windows are more strict about the types here
     cmd = tuple(five.n(arg) for arg in cmd)
     kwargs['env'] = {
         five.n(key): five.n(value)
         for key, value in kwargs.pop('env', {}).items()
     } or None
-    popen_kwargs.update(kwargs)
+    for arg in ('stdin', 'stdout', 'stderr'):
+        kwargs.setdefault(arg, subprocess.PIPE)
+    return cmd, kwargs
+
+
+def cmd_output_b(*cmd, **kwargs):
+    retcode = kwargs.pop('retcode', 0)
+    cmd, kwargs = _cmd_kwargs(*cmd, **kwargs)
 
     try:
         cmd = parse_shebang.normalize_cmd(cmd)
     except parse_shebang.ExecutableNotFoundError as e:
         returncode, stdout_b, stderr_b = e.to_output()
     else:
-        proc = subprocess.Popen(cmd, **popen_kwargs)
+        proc = subprocess.Popen(cmd, **kwargs)
         stdout_b, stderr_b = proc.communicate()
         returncode = proc.returncode
 
@@ -158,6 +157,72 @@ def cmd_output(*cmd, **kwargs):
     return returncode, stdout, stderr
 
 
+if os.name != 'nt':  # pragma: windows no cover
+    from os import openpty
+    import termios
+
+    class Pty(object):
+        def __init__(self):
+            self.r = self.w = None
+
+        def __enter__(self):
+            self.r, self.w = openpty()
+
+            # tty flags normally change \n to \r\n
+            attrs = termios.tcgetattr(self.r)
+            attrs[1] &= ~(termios.ONLCR | termios.OPOST)
+            termios.tcsetattr(self.r, termios.TCSANOW, attrs)
+
+            return self
+
+        def close_w(self):
+            if self.w is not None:
+                os.close(self.w)
+                self.w = None
+
+        def close_r(self):
+            assert self.r is not None
+            os.close(self.r)
+            self.r = None
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            self.close_w()
+            self.close_r()
+
+    def cmd_output_p(*cmd, **kwargs):
+        assert kwargs.pop('retcode') is None
+        assert kwargs['stderr'] == subprocess.STDOUT, kwargs['stderr']
+        cmd, kwargs = _cmd_kwargs(*cmd, **kwargs)
+
+        try:
+            cmd = parse_shebang.normalize_cmd(cmd)
+        except parse_shebang.ExecutableNotFoundError as e:
+            return e.to_output()
+
+        with open(os.devnull) as devnull, Pty() as pty:
+            kwargs.update({'stdin': devnull, 'stdout': pty.w, 'stderr': pty.w})
+            proc = subprocess.Popen(cmd, **kwargs)
+            pty.close_w()
+
+            buf = b''
+            while True:
+                try:
+                    bts = os.read(pty.r, 4096)
+                except OSError as e:
+                    if e.errno == errno.EIO:
+                        bts = b''
+                    else:
+                        raise
+                else:
+                    buf += bts
+                if not bts:
+                    break
+
+        return proc.wait(), buf, None
+else:  # pragma: no cover
+    cmd_output_p = cmd_output_b
+
+
 def rmtree(path):
     """On windows, rmtree fails for readonly dirs."""
     def handle_remove_readonly(func, path, exc):
diff --git a/pre_commit/xargs.py b/pre_commit/xargs.py
@@ -13,6 +13,7 @@
 
 from pre_commit import parse_shebang
 from pre_commit.util import cmd_output_b
+from pre_commit.util import cmd_output_p
 
 
 def _environ_size(_env=None):
@@ -108,9 +109,11 @@ def xargs(cmd, varargs, **kwargs):
     negate: Make nonzero successful and zero a failure
     target_concurrency: Target number of partitions to run concurrently
     """
+    color = kwargs.pop('color', False)
     negate = kwargs.pop('negate', False)
     target_concurrency = kwargs.pop('target_concurrency', 1)
     max_length = kwargs.pop('_max_length', _get_platform_max_length())
+    cmd_fn = cmd_output_p if color else cmd_output_b
     retcode = 0
     stdout = b''
 
@@ -122,7 +125,7 @@ def xargs(cmd, varargs, **kwargs):
     partitions = partition(cmd, varargs, target_concurrency, max_length)
 
     def run_cmd_partition(run_cmd):
-        return cmd_output_b(
+        return cmd_fn(
             *run_cmd, retcode=None, stderr=subprocess.STDOUT, **kwargs
         )
 
diff --git a/testing/resources/stdout_stderr_repo/.pre-commit-hooks.yaml b/testing/resources/stdout_stderr_repo/.pre-commit-hooks.yaml
diff --git a/testing/resources/stdout_stderr_repo/stdout-stderr-entry b/testing/resources/stdout_stderr_repo/stdout-stderr-entry
diff --git a/testing/resources/stdout_stderr_repo/tty-check-entry b/testing/resources/stdout_stderr_repo/tty-check-entry
diff --git a/tests/languages/all_test.py b/tests/languages/all_test.py
diff --git a/tests/parse_shebang_test.py b/tests/parse_shebang_test.py
diff --git a/tests/repository_test.py b/tests/repository_test.py
diff --git a/tests/util_test.py b/tests/util_test.py
diff --git a/tests/xargs_test.py b/tests/xargs_test.py

Original file line number	Diff line number	Diff line change
`@@ -38,16 +38,17 @@`
`38`	`38`	`# version - A version specified in the hook configuration or 'default'.`
`39`	`39`	`# """`
`40`	`40`	`#`
`41`		`-# def run_hook(hook, file_args):`
	`41`	`+# def run_hook(hook, file_args, color):`
`42`	`42`	`# """Runs a hook and returns the returncode and output of running that`
`43`	`43`	`# hook.`
`44`	`44`	`#`
`45`	`45`	`# Args:`
`46`	`46`	# hook - `Hook`
`47`	`47`	`# file_args - The files to be run`
	`48`	`+# color - whether the hook should be given a pty (when supported)`
`48`	`49`	`#`
`49`	`50`	`# Returns:`
`50`		`-# (returncode, stdout, stderr)`
	`51`	`+# (returncode, output)`
`51`	`52`	`# """`
`52`	`53`
`53`	`54`	`languages = {`
Original file line number	Diff line number	Diff line change
`@@ -78,6 +78,6 @@ def install_environment(`
`78`	`78`	`)`
`79`	`79`
`80`	`80`
`81`		`-def run_hook(hook, file_args): # pragma: windows no cover`
	`81`	`+def run_hook(hook, file_args, color): # pragma: windows no cover`
`82`	`82`	`with in_env(hook.prefix, hook.language_version):`
`83`		`- return helpers.run_xargs(hook, helpers.to_cmd(hook), file_args)`
	`83`	`+ return helpers.run_xargs(hook, hook.cmd, file_args, color=color)`