Add --color=[always|never|auto] argument; fix color when piping (#3013)

* Disable spec colorization when redirecting stdout and add command line flag to re-enable
* Add command line `--color` flag to control output colorization
* Add options to `llnl.util.tty.color` to allow color to be auto/always/never
* Add `Spec.cformat()` function to be used when `format()` should have auto-coloring

1 files changed, 54 insertions, 6 deletions

diff --git a/lib/spack/llnl/util/tty/color.py b/lib/spack/llnl/util/tty/color.py
index fc3b697827..a39b5b95f8 100644
--- a/lib/spack/llnl/util/tty/color.py
+++ b/lib/spack/llnl/util/tty/color.py
@@ -80,6 +80,7 @@ To output an @, use '@@'.  To output a } inside braces, use '}}'.
 """
 import re
 import sys
+from contextlib import contextmanager
 
 
 class ColorParseError(Exception):
@@ -107,15 +108,62 @@ colors = {'k': 30, 'K': 90,  # black
 # Regex to be used for color formatting
 color_re = r'@(?:@|\.|([*_])?([a-zA-Z])?(?:{((?:[^}]|}})*)})?)'
 
+# Mapping from color arguments to values for tty.set_color
+color_when_values = {
+    'always': True,
+    'auto': None,
+    'never': False
+}
 
-# Force color even if stdout is not a tty.
-_force_color = False
+# Force color; None: Only color if stdout is a tty
+# True: Always colorize output, False: Never colorize output
+_force_color = None
+
+
+def _color_when_value(when):
+    """Raise a ValueError for an invalid color setting.
+
+    Valid values are 'always', 'never', and 'auto', or equivalently,
+    True, False, and None.
+    """
+    if when in color_when_values:
+        return color_when_values[when]
+    elif when not in color_when_values.values():
+        raise ValueError('Invalid color setting: %s' % when)
+    return when
+
+
+def get_color_when():
+    """Return whether commands should print color or not."""
+    if _force_color is not None:
+        return _force_color
+    return sys.stdout.isatty()
+
+
+def set_color_when(when):
+    """Set when color should be applied.  Options are:
+
+    * True or 'always': always print color
+    * False or 'never': never print color
+    * None or 'auto': only print color if sys.stdout is a tty.
+    """
+    global _force_color
+    _force_color = _color_when_value(when)
+
+
+@contextmanager
+def color_when(value):
+    """Context manager to temporarily use a particular color setting."""
+    old_value = value
+    set_color(value)
+    yield
+    set_color(old_value)
 
 
 class match_to_ansi(object):
 
     def __init__(self, color=True):
-        self.color = color
+        self.color = _color_when_value(color)
 
     def escape(self, s):
         """Returns a TTY escape sequence for a color"""
@@ -166,7 +214,7 @@ def colorize(string, **kwargs):
         color (bool): If False, output will be plain text without control
             codes, for output to non-console devices.
     """
-    color = kwargs.get('color', True)
+    color = _color_when_value(kwargs.get('color', get_color_when()))
     return re.sub(color_re, match_to_ansi(color), string)
 
 
@@ -188,7 +236,7 @@ def cwrite(string, stream=sys.stdout, color=None):
        then it will be set based on stream.isatty().
     """
     if color is None:
-        color = stream.isatty() or _force_color
+        color = get_color_when()
     stream.write(colorize(string, color=color))
 
 
@@ -217,7 +265,7 @@ class ColorStream(object):
             if raw:
                 color = True
             else:
-                color = self._stream.isatty() or _force_color
+                color = get_color_when()
         raw_write(colorize(string, color=color))
 
     def writelines(self, sequence, **kwargs):