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

lib/spack/llnl/util/tty/color.py

@@ -80,6 +80,7 @@
 """
 import re
 import sys
+from contextlib import contextmanager
 
 
 class ColorParseError(Exception):
@@ -107,15 +108,62 @@ def __init__(self, message):
 # 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; None: Only color if stdout is a tty
-_force_color = False
+# 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 @@ def write(self, string, **kwargs):
             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):

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)) +
-        args += [colorize("  @K{%s} " % s.dag_hash(7), color=color) +
+                 s.cformat('$_$@$%@$=') for s in matching_specs]
-                 s.format('$_$@$%@$=', color=color) 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 fmt(s):
                     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 fmt(s):
                 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(

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)

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):

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, "

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,
+    kwargs = {'cover': args.cover,
-              'cover': args.cover,
               'format': name_fmt + '$@$%@+$+$=',
               'hashes': args.long or args.very_long,
               'hashlen': None if args.very_long else 7,

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(
+            tty.error("Will not uninstall %s" % spec.cformat("$_$@$%@$/"))
-                      spec.format("$_$@$%@$/", color=True)))
             print('')
             print('The following packages depend on it:')
             spack.cmd.display_specs(lst, **display_args)

lib/spack/spack/database.py

@@ -288,7 +288,7 @@ def _assign_dependencies(self, hash_key, installs, data):
                 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 @@ def _read_suppress_error():
                     # 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 @@ def _read_suppress_error():
                     # 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:

lib/spack/spack/directory_layout.py

@@ -168,7 +168,9 @@ def __init__(self, root, **kwargs):
         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(

lib/spack/spack/main.py

@@ -58,7 +58,7 @@
 
 # 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.

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"
+            tty.warn(
-                     % spec.format('$_$@'), e.message)
+                "Error while fetching %s" % spec.cformat('$_$@'), e.message)
         categories['error'].append(spec)
 
 

lib/spack/spack/package.py

@@ -905,7 +905,7 @@ def do_fetch(self, mirror_only=False):
         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 @@ def do_activate(self, force=False):
         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" %
+        tty.msg(
-                (self.spec.short_spec, self.extendee_spec.format("$_$@$+$%@")))
+            "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 @@ def do_deactivate(self, **kwargs):
             spack.store.layout.remove_extension(
                 self.extendee_spec, self.spec)
 
-        tty.msg("Deactivated extension %s for %s" %
+        tty.msg(
-                (self.spec.short_spec, self.extendee_spec.format("$_$@$+$%@")))
+            "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.

lib/spack/spack/spec.py

@@ -108,6 +108,10 @@
 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.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 @@ def short_spec(self):
 
     @property
     def cshort_spec(self):
-        """Returns a version of the spec with the dependencies hashed
+        """Returns an auto-colorized version of ``self.short_spec``."""
-           instead of completely enumerated."""
+        return self.cformat('$_$@$%@$+$=$/')
-        return self.format('$_$@$%@$+$=$/', color=True)
 
     @property
     def prefix(self):
@@ -2852,6 +2852,12 @@ def write(s, c):
         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 @@ def _installed_explicitly(self):
     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)

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