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

14 files changed, 105 insertions, 45 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):
diff --git a/lib/spack/spack/cmd/__init__.py b/lib/spack/spack/cmd/__init__.py
index edeaa677de..b9b145d0b5 100644
--- a/lib/spack/spack/cmd/__init__.py
+++ b/lib/spack/spack/cmd/__init__.py
@@ -154,9 +154,8 @@ def disambiguate_spec(spec):
     elif len(matching_specs) > 1:
         args = ["%s matches multiple packages." % spec,
                 "Matching packages:"]
-        color = sys.stdout.isatty()
-        args += [colorize("  @K{%s} " % s.dag_hash(7), color=color) +
-                 s.format('$_$@$%@$=', color=color) for s in matching_specs]
+        args += [colorize("  @K{%s} " % s.dag_hash(7)) +
+                 s.cformat('$_$@$%@$=') for s in matching_specs]
         args += ["Use a more specific spec."]
         tty.die(*args)
 
@@ -200,7 +199,7 @@ def display_specs(specs, **kwargs):
         specs = index[(architecture, compiler)]
         specs.sort()
 
-        abbreviated = [s.format(format_string, color=True) for s in specs]
+        abbreviated = [s.cformat(format_string) for s in specs]
         if mode == 'paths':
             # Print one spec per line along with prefix path
             width = max(len(s) for s in abbreviated)
@@ -215,7 +214,6 @@ def display_specs(specs, **kwargs):
             for spec in specs:
                 print(spec.tree(
                     format=format_string,
-                    color=True,
                     indent=4,
                     prefix=(lambda s: gray_hash(s, hlen)) if hashes else None))
 
@@ -227,7 +225,7 @@ def display_specs(specs, **kwargs):
                     string = ""
                     if hashes:
                         string += gray_hash(s, hlen) + ' '
-                    string += s.format('$-%s$@%s' % (nfmt, vfmt), color=True)
+                    string += s.cformat('$-%s$@%s' % (nfmt, vfmt))
 
                     return string
 
@@ -237,7 +235,7 @@ def display_specs(specs, **kwargs):
                 for spec in specs:
                     # Print the hash if necessary
                     hsh = gray_hash(spec, hlen) + ' ' if hashes else ''
-                    print(hsh + spec.format(format_string, color=True) + '\n')
+                    print(hsh + spec.cformat(format_string) + '\n')
 
         else:
             raise ValueError(
diff --git a/lib/spack/spack/cmd/dependents.py b/lib/spack/spack/cmd/dependents.py
index c983dd79ce..e8f8fc0d0b 100644
--- a/lib/spack/spack/cmd/dependents.py
+++ b/lib/spack/spack/cmd/dependents.py
@@ -47,7 +47,7 @@ def dependents(parser, args):
         tty.die("spack dependents takes only one spec.")
     spec = spack.cmd.disambiguate_spec(specs[0])
 
-    tty.msg("Dependents of %s" % spec.format('$_$@$%@$/', color=True))
+    tty.msg("Dependents of %s" % spec.cformat('$_$@$%@$/'))
     deps = spack.store.db.installed_dependents(spec)
     if deps:
         spack.cmd.display_specs(deps)
diff --git a/lib/spack/spack/cmd/mirror.py b/lib/spack/spack/cmd/mirror.py
index bfc36c4107..b89ac4121d 100644
--- a/lib/spack/spack/cmd/mirror.py
+++ b/lib/spack/spack/cmd/mirror.py
@@ -213,7 +213,7 @@ def mirror_create(args):
         "  %-4d failed to fetch." % e)
     if error:
         tty.error("Failed downloads:")
-        colify(s.format("$_$@") for s in error)
+        colify(s.cformat("$_$@") for s in error)
 
 
 def mirror(parser, args):
diff --git a/lib/spack/spack/cmd/module.py b/lib/spack/spack/cmd/module.py
index ed20ad4029..d59d4433b7 100644
--- a/lib/spack/spack/cmd/module.py
+++ b/lib/spack/spack/cmd/module.py
@@ -236,7 +236,8 @@ def refresh(mtype, specs, args):
             if len(writer_list) > 1:
                 message += '\nfile: {0}\n'.format(filename)
                 for x in writer_list:
-                    message += 'spec: {0}\n'.format(x.spec.format(color=True))
+                    message += 'spec: {0}\n'.format(x.spec.format())
+
         tty.error(message)
         tty.error('Operation aborted')
         raise SystemExit(1)
@@ -269,7 +270,7 @@ def module(parser, args):
                    "and this is not allowed in this context")
         tty.error(message.format(query=constraint))
         for s in specs:
-            sys.stderr.write(s.format(color=True) + '\n')
+            sys.stderr.write(s.cformat() + '\n')
         raise SystemExit(1)
     except NoMatch:
         message = ("the constraint '{query}' matches no package, "
diff --git a/lib/spack/spack/cmd/spec.py b/lib/spack/spack/cmd/spec.py
index ee1ec882e2..b4740b4ece 100644
--- a/lib/spack/spack/cmd/spec.py
+++ b/lib/spack/spack/cmd/spec.py
@@ -60,8 +60,7 @@ def setup_parser(subparser):
 
 def spec(parser, args):
     name_fmt = '$.' if args.namespaces else '$_'
-    kwargs = {'color': True,
-              'cover': args.cover,
+    kwargs = {'cover': args.cover,
               'format': name_fmt + '$@$%@+$+$=',
               'hashes': args.long or args.very_long,
               'hashlen': None if args.very_long else 7,
diff --git a/lib/spack/spack/cmd/uninstall.py b/lib/spack/spack/cmd/uninstall.py
index 499521ede0..d4e88288d2 100644
--- a/lib/spack/spack/cmd/uninstall.py
+++ b/lib/spack/spack/cmd/uninstall.py
@@ -180,8 +180,7 @@ def get_uninstall_list(args):
     has_error = False
     if dependent_list and not args.dependents and not args.force:
         for spec, lst in dependent_list.items():
-            tty.error('Will not uninstall {0}'.format(
-                      spec.format("$_$@$%@$/", color=True)))
+            tty.error("Will not uninstall %s" % spec.cformat("$_$@$%@$/"))
             print('')
             print('The following packages depend on it:')
             spack.cmd.display_specs(lst, **display_args)
diff --git a/lib/spack/spack/database.py b/lib/spack/spack/database.py
index 75862ff87c..569de40a94 100644
--- a/lib/spack/spack/database.py
+++ b/lib/spack/spack/database.py
@@ -288,7 +288,7 @@ class Database(object):
                 if dhash not in data:
                     tty.warn("Missing dependency not in database: ",
                              "%s needs %s-%s" % (
-                                 spec.format('$_$/'), dname, dhash[:7]))
+                                 spec.cformat('$_$/'), dname, dhash[:7]))
                     continue
 
                 child = data[dhash].spec
@@ -440,8 +440,7 @@ class Database(object):
                     # just to be conservative in case a command like
                     # "autoremove" is run by the user after a reindex.
                     tty.debug(
-                        'RECONSTRUCTING FROM SPEC.YAML: {0}'.format(spec)
-                    )
+                        'RECONSTRUCTING FROM SPEC.YAML: {0}'.format(spec))
                     explicit = True
                     if old_data is not None:
                         old_info = old_data.get(spec.dag_hash())
@@ -467,8 +466,7 @@ class Database(object):
                     # installed compilers or externally installed
                     # applications.
                     tty.debug(
-                        'RECONSTRUCTING FROM OLD DB: {0}'.format(entry.spec)
-                    )
+                        'RECONSTRUCTING FROM OLD DB: {0}'.format(entry.spec))
                     try:
                         layout = spack.store.layout
                         if entry.spec.external:
diff --git a/lib/spack/spack/directory_layout.py b/lib/spack/spack/directory_layout.py
index 5e9341ced9..5b430224d6 100644
--- a/lib/spack/spack/directory_layout.py
+++ b/lib/spack/spack/directory_layout.py
@@ -168,7 +168,9 @@ class YamlDirectoryLayout(DirectoryLayout):
         self.metadata_dir   = kwargs.get('metadata_dir', '.spack')
         self.hash_len       = kwargs.get('hash_len')
         self.path_scheme    = kwargs.get('path_scheme') or (
-            "${ARCHITECTURE}/${COMPILERNAME}-${COMPILERVER}/${PACKAGE}-${VERSION}-${HASH}")  # NOQA: E501
+            "${ARCHITECTURE}/"
+            "${COMPILERNAME}-${COMPILERVER}/"
+            "${PACKAGE}-${VERSION}-${HASH}")
         if self.hash_len is not None:
             if re.search('\${HASH:\d+}', self.path_scheme):
                 raise InvalidDirectoryLayoutParametersError(
diff --git a/lib/spack/spack/main.py b/lib/spack/spack/main.py
index 2d5648f13e..ee74c915df 100644
--- a/lib/spack/spack/main.py
+++ b/lib/spack/spack/main.py
@@ -58,7 +58,7 @@ intro_by_level = {
 
 # control top-level spack options shown in basic vs. advanced help
 options_by_level = {
-    'short': 'hkV',
+    'short': ['h', 'k', 'V', 'color'],
     'long': 'all'
 }
 
@@ -280,6 +280,9 @@ def make_argument_parser():
 
     parser.add_argument('-h', '--help', action='store_true',
                         help="show this help message and exit")
+    parser.add_argument('--color', action='store', default='auto',
+                        choices=('always', 'never', 'auto'),
+                        help="when to colorize output; default is auto")
     parser.add_argument('-d', '--debug', action='store_true',
                         help="write out debug logs during compile")
     parser.add_argument('-D', '--pdb', action='store_true',
@@ -325,6 +328,9 @@ def setup_main_options(args):
         tty.warn("You asked for --insecure. Will NOT check SSL certificates.")
         spack.insecure = True
 
+    # when to use color (takes always, auto, or never)
+    tty.color.set_color_when(args.color)
+
 
 def allows_unknown_args(command):
     """Implements really simple argument injection for unknown arguments.
diff --git a/lib/spack/spack/mirror.py b/lib/spack/spack/mirror.py
index 9b9f816db4..bfa7b858d2 100644
--- a/lib/spack/spack/mirror.py
+++ b/lib/spack/spack/mirror.py
@@ -224,14 +224,14 @@ def add_single_spec(spec, mirror_root, categories, **kwargs):
                     # create a subdirectory for the current package@version
                     archive_path = os.path.abspath(join_path(
                         mirror_root, mirror_archive_path(spec, fetcher)))
-                    name = spec.format("$_$@")
+                    name = spec.cformat("$_$@")
                 else:
                     resource = stage.resource
                     archive_path = os.path.abspath(join_path(
                         mirror_root,
                         mirror_archive_path(spec, fetcher, resource.name)))
                     name = "{resource} ({pkg}).".format(
-                        resource=resource.name, pkg=spec.format("$_$@"))
+                        resource=resource.name, pkg=spec.cformat("$_$@"))
                 subdir = os.path.dirname(archive_path)
                 mkdirp(subdir)
 
@@ -258,8 +258,8 @@ def add_single_spec(spec, mirror_root, categories, **kwargs):
         if spack.debug:
             sys.excepthook(*sys.exc_info())
         else:
-            tty.warn("Error while fetching %s"
-                     % spec.format('$_$@'), e.message)
+            tty.warn(
+                "Error while fetching %s" % spec.cformat('$_$@'), e.message)
         categories['error'].append(spec)
 
 
diff --git a/lib/spack/spack/package.py b/lib/spack/spack/package.py
index 1223fce178..5729836873 100644
--- a/lib/spack/spack/package.py
+++ b/lib/spack/spack/package.py
@@ -905,7 +905,7 @@ class PackageBase(with_metaclass(PackageMeta, object)):
         start_time = time.time()
         if spack.do_checksum and self.version not in self.versions:
             tty.warn("There is no checksum on file to fetch %s safely." %
-                     self.spec.format('$_$@'))
+                     self.spec.cformat('$_$@'))
 
             # Ask the user whether to skip the checksum if we're
             # interactive, but just fail if non-interactive.
@@ -1649,8 +1649,9 @@ class PackageBase(with_metaclass(PackageMeta, object)):
         self.extendee_spec.package.activate(self, **self.extendee_args)
 
         spack.store.layout.add_extension(self.extendee_spec, self.spec)
-        tty.msg("Activated extension %s for %s" %
-                (self.spec.short_spec, self.extendee_spec.format("$_$@$+$%@")))
+        tty.msg(
+            "Activated extension %s for %s" %
+            (self.spec.short_spec, self.extendee_spec.cformat("$_$@$+$%@")))
 
     def dependency_activations(self):
         return (spec for spec in self.spec.traverse(root=False, deptype='run')
@@ -1708,8 +1709,9 @@ class PackageBase(with_metaclass(PackageMeta, object)):
             spack.store.layout.remove_extension(
                 self.extendee_spec, self.spec)
 
-        tty.msg("Deactivated extension %s for %s" %
-                (self.spec.short_spec, self.extendee_spec.format("$_$@$+$%@")))
+        tty.msg(
+            "Deactivated extension %s for %s" %
+            (self.spec.short_spec, self.extendee_spec.cformat("$_$@$+$%@")))
 
     def deactivate(self, extension, **kwargs):
         """Unlinks all files from extension out of this package's install dir.
diff --git a/lib/spack/spack/spec.py b/lib/spack/spack/spec.py
index db8dcf61a8..992930da65 100644
--- a/lib/spack/spack/spec.py
+++ b/lib/spack/spack/spec.py
@@ -108,6 +108,10 @@ from six import StringIO
 from six import string_types
 from six import iteritems
 
+from llnl.util.filesystem import find_headers, find_libraries, is_exe
+from llnl.util.lang import *
+from llnl.util.tty.color import *
+
 import spack
 import spack.architecture
 import spack.compilers as compilers
@@ -117,9 +121,6 @@ import spack.store
 import spack.util.spack_json as sjson
 import spack.util.spack_yaml as syaml
 
-from llnl.util.filesystem import find_headers, find_libraries, is_exe
-from llnl.util.lang import *
-from llnl.util.tty.color import *
 from spack.util.module_cmd import get_path_from_module, load_module
 from spack.error import SpecError, UnsatisfiableSpecError
 from spack.provider_index import ProviderIndex
@@ -1363,9 +1364,8 @@ class Spec(object):
 
     @property
     def cshort_spec(self):
-        """Returns a version of the spec with the dependencies hashed
-           instead of completely enumerated."""
-        return self.format('$_$@$%@$+$=$/', color=True)
+        """Returns an auto-colorized version of ``self.short_spec``."""
+        return self.cformat('$_$@$%@$+$=$/')
 
     @property
     def prefix(self):
@@ -2852,6 +2852,12 @@ class Spec(object):
         result = out.getvalue()
         return result
 
+    def cformat(self, *args, **kwargs):
+        """Same as format, but color defaults to auto instead of False."""
+        kwargs = kwargs.copy()
+        kwargs.setdefault('color', None)
+        return self.format(*args, **kwargs)
+
     def dep_string(self):
         return ''.join("^" + dep.format() for dep in self.sorted_deps())
 
@@ -2882,7 +2888,7 @@ class Spec(object):
     def tree(self, **kwargs):
         """Prints out this spec and its dependencies, tree-formatted
            with indentation."""
-        color = kwargs.pop('color', False)
+        color = kwargs.pop('color', get_color_when())
         depth = kwargs.pop('depth', False)
         hashes = kwargs.pop('hashes', False)
         hlen = kwargs.pop('hashlen', None)
diff --git a/share/spack/spack-completion.bash b/share/spack/spack-completion.bash
index d9ae65b8b5..e60d587c18 100755
--- a/share/spack/spack-completion.bash
+++ b/share/spack/spack-completion.bash
@@ -115,7 +115,8 @@ function _spack {
     if $list_options
     then
         compgen -W "-h --help -d --debug -D --pdb -k --insecure -m --mock -p
-                    --profile -v --verbose -s --stacktrace -V --version" -- "$cur"
+                    --profile -v --verbose -s --stacktrace -V --version
+                    --color --color=always --color=auto --color=never" -- "$cur"
     else
         compgen -W "$(_subcommands)" -- "$cur"
     fi