Merge pull request #64 from pre-commit/color_63

asottile · asottile · commit 749615118e16 · 2014-04-05T19:34:59.000-07:00
Wire in color for pre-commit
diff --git a/pre_commit/color.py b/pre_commit/color.py
@@ -0,0 +1,38 @@
+
+import sys
+
+RED = '\033[41m'
+GREEN = '\033[42m'
+NORMAL = '\033[0m'
+
+
+class InvalidColorSetting(ValueError): pass
+
+
+def format_color(text, color, use_color):
+    """Format text with color.
+
+    Args:
+        text - Text to be formatted with color if `use_color`
+        color - The color start string
+        use_color - Whether or not to color
+    """
+    if not use_color:
+        return text
+    else:
+        return u'{0}{1}{2}'.format(color, text, NORMAL)
+
+
+def use_color(setting):
+    """Choose whether to use color based on the command argument.
+
+    Args:
+        setting - Either `auto`, `always`, or `never`
+    """
+    if setting not in ('auto', 'always', 'never'):
+        raise InvalidColorSetting(setting)
+
+    return (
+        setting == 'always' or
+        (setting == 'auto' and sys.stdout.isatty())
+    )
diff --git a/pre_commit/run.py b/pre_commit/run.py
@@ -5,22 +5,20 @@
 import subprocess
 import sys
 
+from pre_commit import color
 from pre_commit import commands
 from pre_commit import git
 from pre_commit.runner import Runner
 from pre_commit.util import entry
 
 
-RED = '\033[41m'
-GREEN = '\033[42m'
-NORMAL = '\033[0m'
 COLS = int(subprocess.Popen(['tput', 'cols'], stdout=subprocess.PIPE).communicate()[0])
 
 PASS_FAIL_LENGTH = 6
 
 
-def _run_single_hook(runner, repository, hook_id, all_files=False, verbose=False):
-    if all_files:
+def _run_single_hook(runner, repository, hook_id, args):
+    if args.all_files:
         get_filenames = git.get_all_files_matching
     else:
         get_filenames = git.get_staged_files_matching
@@ -46,54 +44,49 @@ def _run_single_hook(runner, repository, hook_id, all_files=False, verbose=False
     output = '\n'.join([stdout, stderr]).strip()
     if retcode != repository.hooks[hook_id]['expected_return_value']:
         retcode = 1
-        color = RED
+        print_color = color.RED
         pass_fail = 'Failed'
     else:
         retcode = 0
-        color = GREEN
+        print_color = color.GREEN
         pass_fail = 'Passed'
 
 
-    print('{0}{1}{2}'.format(color, pass_fail, NORMAL))
+    print(color.format_color(pass_fail, print_color, args.color))
 
-    if output and (retcode or verbose):
+    if output and (retcode or args.verbose):
         print('\n' + output)
 
     return retcode
 
 
-def run_hooks(runner, all_files=False, verbose=False):
+def run_hooks(runner, args):
     """Actually run the hooks."""
     retval = 0
 
     for repo in runner.repositories:
         for hook_id in repo.hooks:
-            retval |= _run_single_hook(
-                runner,
-                repo,
-                hook_id,
-                all_files=all_files,
-                verbose=verbose,
-            )
+            retval |= _run_single_hook(runner, repo, hook_id, args)
 
     return retval
 
 
-def run_single_hook(runner, hook_id, all_files=False, verbose=False):
+def run_single_hook(runner, hook_id, args):
     for repo in runner.repositories:
         if hook_id in repo.hooks:
-            return _run_single_hook(
-                runner,
-                repo,
-                hook_id,
-                all_files=all_files,
-                verbose=verbose,
-            )
+            return _run_single_hook(runner, repo, hook_id, args)
     else:
         print('No hook with id `{0}`'.format(hook_id))
         return 1
 
 
+def _run(runner, args):
+    if args.hook:
+        return run_single_hook(runner, args.hook, args)
+    else:
+        return run_hooks(runner, args)
+
+
 @entry
 def run(argv):
     parser = argparse.ArgumentParser()
@@ -115,6 +108,10 @@ def run(argv):
         help='Run on all the files in the repo.',
     )
     run.add_argument('--verbose', '-v', action='store_true', default=False)
+    run.add_argument(
+        '--color', default='auto', type=color.use_color,
+        help='Whether to use color in output.  Defaults to `auto`',
+    )
 
     help = subparsers.add_parser('help', help='Show help for a specific command.')
     help.add_argument('help_cmd', nargs='?', help='Command to show help for.')
@@ -135,17 +132,7 @@ def run(argv):
     elif args.command == 'autoupdate':
         return commands.autoupdate(runner)
     elif args.command == 'run':
-        if args.hook:
-            return run_single_hook(
-                runner,
-                args.hook,
-                all_files=args.all_files,
-                verbose=args.verbose,
-            )
-        else:
-            return run_hooks(
-                runner, all_files=args.all_files, verbose=args.verbose,
-            )
+        return _run(runner, args)
     elif args.command == 'help':
         if args.help_cmd:
             parser.parse_args([args.help_cmd, '--help'])
diff --git a/tests/color_test.py b/tests/color_test.py
@@ -0,0 +1,41 @@
+
+import mock
+import pytest
+import sys
+
+from pre_commit.color import format_color
+from pre_commit.color import GREEN
+from pre_commit.color import InvalidColorSetting
+from pre_commit.color import use_color
+
+
+@pytest.mark.parametrize(('in_text', 'in_color', 'in_use_color', 'expected'), (
+    ('foo', GREEN, True, '{0}foo\033[0m'.format(GREEN)),
+    ('foo', GREEN, False, 'foo'),
+))
+def test_format_color(in_text, in_color, in_use_color, expected):
+    ret = format_color(in_text, in_color, in_use_color)
+    assert ret == expected
+
+
+def test_use_color_never():
+    assert use_color('never') is False
+
+
+def test_use_color_always():
+    assert use_color('always') is True
+
+
+def test_use_color_no_tty():
+    with mock.patch.object(sys.stdout, 'isatty', return_value=False):
+        assert use_color('auto') is False
+
+
+def test_use_color_tty():
+    with mock.patch.object(sys.stdout, 'isatty', return_value=True):
+        assert use_color('auto') is True
+
+
+def test_use_color_raises_if_given_shenanigans():
+    with pytest.raises(InvalidColorSetting):
+        use_color('herpaderp')
