Add color utility.

asottile · asottile · commit 6a1f945e3105 · 2014-04-05T19:12:30.000-07:00
diff --git a/pre_commit/color.py b/pre_commit/color.py
@@ -0,0 +1,32 @@
+
+import sys
+
+RED = '\033[41m'
+GREEN = '\033[42m'
+NORMAL = '\033[0m'
+
+
+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`
+    """
+    return (
+        setting == 'always' or
+        (setting == 'auto' and sys.stdout.isatty())
+    )
diff --git a/tests/color_test.py b/tests/color_test.py
@@ -0,0 +1,35 @@
+
+import mock
+import pytest
+import sys
+
+from pre_commit.color import format_color
+from pre_commit.color import GREEN
+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
