refactor color utils

1 files changed, 20 insertions, 23 deletions

diff --git a/psutil/tests/runner.py b/psutil/tests/runner.py
index f8601bad..aa81378a 100755
--- a/psutil/tests/runner.py
+++ b/psutil/tests/runner.py
@@ -31,15 +31,6 @@ from psutil.tests import TOX
 HERE = os.path.abspath(os.path.dirname(__file__))
 VERBOSITY = 1 if TOX else 2
 FAILED_TESTS_FNAME = '.failed-tests.txt'
-if os.name == 'posix':
-    GREEN = 1
-    RED = 2
-    BROWN = 94
-else:
-    GREEN = 2
-    RED = 4
-    BROWN = 6
-    DEFAULT_COLOR = 7
 
 
 def term_supports_colors(file=sys.stdout):
@@ -56,17 +47,16 @@ def term_supports_colors(file=sys.stdout):
         return True
 
 
-def hilite(s, color, bold=False):
+def hilite(s, color="green", bold=False):
     """Return an highlighted version of 'string'."""
     attr = []
-    if color == GREEN:
-        attr.append('32')
-    elif color == RED:
-        attr.append('91')
-    elif color == BROWN:
-        attr.append('33')
-    else:
-        raise ValueError("unrecognized color")
+    colors = dict(green='32', red='91', brown='33')
+    try:
+        color = colors[color]
+    except KeyError:
+        raise ValueError("invalid color %r; choose between %r" % (
+            list(colors.keys())))
+    attr.append(color)
     if bold:
         attr.append('1')
     return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), s)
@@ -83,6 +73,13 @@ def _stderr_handle():
 
 
 def win_colorprint(printer, s, color, bold=False):
+    DEFAULT_COLOR = 7
+    colors = dict(green=2, red=4, brown=6)
+    try:
+        color = colors[color]
+    except KeyError:
+        raise ValueError("invalid color %r; choose between %r" % (
+            list(colors.keys())))
     if bold and color <= 7:
         color += 8
     handle = _stderr_handle()
@@ -104,22 +101,22 @@ class ColouredResult(TextTestResult):
 
     def addSuccess(self, test):
         TestResult.addSuccess(self, test)
-        self._color_print("OK", GREEN)
+        self._color_print("OK", "green")
 
     def addError(self, test, err):
         TestResult.addError(self, test, err)
-        self._color_print("ERROR", RED, bold=True)
+        self._color_print("ERROR", "red", bold=True)
 
     def addFailure(self, test, err):
         TestResult.addFailure(self, test, err)
-        self._color_print("FAIL", RED)
+        self._color_print("FAIL", "red")
 
     def addSkip(self, test, reason):
         TestResult.addSkip(self, test, reason)
-        self._color_print("skipped: %s" % reason, BROWN)
+        self._color_print("skipped: %s" % reason, "brown")
 
     def printErrorList(self, flavour, errors):
-        flavour = hilite(flavour, RED, bold=flavour == 'ERROR')
+        flavour = hilite(flavour, "red", bold=flavour == 'ERROR')
         TextTestResult.printErrorList(self, flavour, errors)