Color tweaks for find.

2 files changed, 47 insertions, 22 deletions

diff --git a/lib/spack/spack/cmd/find.py b/lib/spack/spack/cmd/find.py
index 9ad8df9489..0eeeeb10a2 100644
--- a/lib/spack/spack/cmd/find.py
+++ b/lib/spack/spack/cmd/find.py
@@ -1,9 +1,12 @@
 import collections
 import argparse
+from StringIO import StringIO
 
 import spack
+import spack.spec
 import spack.packages as packages
 import spack.colify
+from spack.color import *
 from spack.colify import colify
 
 description ="Find installed spack packages"
@@ -21,7 +24,7 @@ def setup_parser(subparser):
 
 
 # TODO: move this and colify to tty.
-def hline(label, char):
+def hline(label, char, color=''):
     max_width = 64
     cols, rows = spack.colify.get_terminal_size()
     if not cols:
@@ -31,9 +34,18 @@ def hline(label, char):
     cols = min(max_width, cols)
 
     label = str(label)
-    out =  char * 2 + " " + label + " "
-    out += (cols - len(out)) * char
-    return out
+    prefix = char * 2 + " " + label + " "
+    suffix = (cols - len(prefix)) * char
+
+    out = StringIO()
+    if color:
+        prefix = char * 2 + " " + color + cescape(label) + "@. "
+        cwrite(prefix, stream=out, color=True)
+    else:
+        out.write(prefix)
+    out.write(suffix)
+
+    return out.getvalue()
 
 
 def find(parser, args):
@@ -56,14 +68,14 @@ def find(parser, args):
 
     # Traverse the index and print out each package
     for architecture in index:
-        print hline(architecture, "=")
+        print hline(architecture, "=", spack.spec.architecture_color)
         for compiler in index[architecture]:
-            print hline(compiler, "-")
+            print hline(compiler, "-", spack.spec.compiler_color)
 
             specs = index[architecture][compiler]
             specs.sort()
 
-            abbreviated = [s.format('$_$@$+$#') for s in specs]
+            abbreviated = [s.format('$_$@$+$#', color=True) for s in specs]
 
             if args.paths:
                 # Print one spec per line along with prefix path
@@ -76,7 +88,7 @@ def find(parser, args):
 
             elif args.full_specs:
                 for spec in specs:
-                    print spec.tree(indent=4, format='$_$@$+'),
+                    print spec.tree(indent=4, format='$_$@$+', color=True),
             else:
                 for abbrv in abbreviated:
                     print "    %s" % abbrv
diff --git a/lib/spack/spack/spec.py b/lib/spack/spack/spec.py
index 5d166861a1..efe5df51d5 100644
--- a/lib/spack/spack/spec.py
+++ b/lib/spack/spack/spec.py
@@ -83,15 +83,23 @@ from spack.util.lang import *
 from spack.util.string import *
 
 
+# Convenient names for color formats so that other things can use them
+compiler_color         = '@g'
+version_color          = '@c'
+architecture_color     = '@m'
+enabled_variant_color  = '@B'
+disabled_variant_color = '@r'
+dependency_color       = '@.'
+
 """This map determines the coloring of specs when using color output.
    We make the fields different colors to enhance readability.
    See spack.color for descriptions of the color codes. """
-color_formats = {'%' : '@g',   # compiler
-                 '@' : '@c',   # version
-                 '=' : '@m',   # architecture
-                 '+' : '@B',   # enable variant
-                 '~' : '@r',   # disable variant
-                 '^' : '@.'}   # dependency
+color_formats = {'%' : compiler_color,
+                 '@' : version_color,
+                 '=' : architecture_color,
+                 '+' : enabled_variant_color,
+                 '~' : disabled_variant_color,
+                 '^' : dependency_color }
 
 """Regex used for splitting by spec field separators."""
 separators = '[%s]' % ''.join(color_formats.keys())
@@ -823,27 +831,34 @@ class Spec(object):
                      of the package, but no dependencies, arch, or compiler.
            """
         color = kwargs.get('color', False)
-
         length = len(format_string)
         out = StringIO()
         escape = compiler = False
+
+        def write(s, c):
+            if color:
+                f = color_formats[c] + cescape(s) + '@.'
+                cwrite(f, stream=out, color=color)
+            else:
+                out.write(s)
+
         for i, c in enumerate(format_string):
             if escape:
                 if c == '_':
                     out.write(self.name)
                 elif c == '@':
                     if self.versions and self.versions != VersionList([':']):
-                        out.write(c + str(self.versions))
+                        write(c + str(self.versions), c)
                 elif c == '%':
                     if self.compiler:
-                        out.write(c + str(self.compiler.name))
+                        write(c + str(self.compiler.name), c)
                     compiler = True
                 elif c == '+':
                     if self.variants:
-                        out.write(str(self.variants))
+                        write(str(self.variants), c)
                 elif c == '=':
                     if self.architecture:
-                        out.write(c + str(self.architecture))
+                        write(c + str(self.architecture), c)
                 elif c == '#':
                     if self.dependencies:
                         out.write('-' + self.dependencies.sha1()[:6])
@@ -854,7 +869,7 @@ class Spec(object):
             elif compiler:
                 if c == '@':
                     if self.compiler and self.compiler.versions:
-                        out.write(c + str(self.compiler.versions))
+                        write(c + str(self.compiler.versions), '%')
                 elif c == '$':
                     escape = True
                 else:
@@ -870,8 +885,6 @@ class Spec(object):
                 out.write(c)
 
         result = out.getvalue()
-        if color:
-            result = colorize_spec(result)
         return result