From 73e0111e383945abfa4251e84549a6cc11651a7e Mon Sep 17 00:00:00 2001
From: Anthony Sottile <asottile@umich.edu>
Date: Sat, 5 Apr 2014 19:31:31 -0700
Subject: [PATCH] Wire in color for pre-commit.  Closes #63.

---
 pre_commit/color.py |  6 ++++++
 pre_commit/run.py   | 14 ++++++++------
 tests/color_test.py |  6 ++++++
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/pre_commit/color.py b/pre_commit/color.py
index d3875c39..641e719e 100644
--- a/pre_commit/color.py
+++ b/pre_commit/color.py
@@ -6,6 +6,9 @@ GREEN = '\033[42m'
 NORMAL = '\033[0m'
 
 
+class InvalidColorSetting(ValueError): pass
+
+
 def format_color(text, color, use_color):
     """Format text with color.
 
@@ -26,6 +29,9 @@ def use_color(setting):
     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
index 20853e0a..7f1cb6b5 100644
--- a/pre_commit/run.py
+++ b/pre_commit/run.py
@@ -5,15 +5,13 @@ import argparse
 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
@@ -46,15 +44,15 @@ def _run_single_hook(runner, repository, hook_id, args):
     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 args.verbose):
         print('\n' + output)
@@ -110,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.')
diff --git a/tests/color_test.py b/tests/color_test.py
index 930861f9..44d9b170 100644
--- a/tests/color_test.py
+++ b/tests/color_test.py
@@ -5,6 +5,7 @@ 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
 
 
@@ -33,3 +34,8 @@ def test_use_color_no_tty():
 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')