concretizer: colorize ASP output

lib/spack/spack/solver/asp.py

@@ -13,6 +13,7 @@
 from six import string_types
 
 import llnl.util.tty as tty
+import llnl.util.tty.color as color
 
 import spack
 import spack.cmd
@@ -386,6 +387,37 @@ def __init__(self, asp):
         self.answers = []
 
 
+def highlight(string):
+    """Syntax highlighting for ASP programs"""
+    # variables
+    string = re.sub(r'\b([A-Z])\b', r'@y{\1}', string)
+
+    # implications
+    string = re.sub(r':-', r'@*G{:-}', string)
+
+    # final periods
+    string = re.sub(r'^([^%].*)\.$', r'\1@*G{.}', string, flags=re.MULTILINE)
+
+    # directives
+    string = re.sub(
+        r'(#\w*)( (?:\w*)?)((?:/\d+)?)', r'@*B{\1}@c{\2}\3', string)
+
+    # functions
+    string = re.sub(r'(\w[\w-]+)\(([^)]*)\)', r'@C{\1}@w{(}\2@w{)}', string)
+
+    # comments
+    string = re.sub(r'(%.*)$', r'@K\1@.', string, flags=re.MULTILINE)
+
+    # strings
+    string = re.sub(r'("[^"]*")', r'@m{\1}', string)
+
+    # result
+    string = re.sub(r'\bUNSATISFIABLE', "@R{UNSATISFIABLE}", string)
+    string = re.sub(r'\bINCONSISTENT', "@R{INCONSISTENT}", string)
+    string = re.sub(r'\bSATISFIABLE', "@G{SATISFIABLE}", string)
+
+    return string
+
 #
 # These are handwritten parts for the Spack ASP model.
 #
@@ -400,6 +432,9 @@ def solve(specs, dump=None, models=1):
     clingo = which('clingo', required=True)
     parser = ResultParser()
 
+    def colorize(string):
+        color.cprint(highlight(color.cescape(string)))
+
     with tempfile.TemporaryFile("w+") as program:
         generator = AspGenerator(program)
         generator.generate_asp_program(specs)
@@ -409,8 +444,9 @@ def solve(specs, dump=None, models=1):
         program.seek(0)
 
         if 'asp' in dump:
+            if sys.stdout.isatty():
                 tty.msg('ASP program:')
-            sys.stdout.write(result.asp)
+            colorize(result.asp)
 
         with tempfile.TemporaryFile("w+") as output:
             with tempfile.TemporaryFile() as warnings:
@@ -427,9 +463,11 @@ def solve(specs, dump=None, models=1):
                 # dump any warnings generated by the solver
                 if 'warnings' in dump:
                     if result.warnings:
+                        if sys.stdout.isatty():
                             tty.msg('Clingo gave the following warnings:')
-                        sys.stdout.write(result.warnings)
+                        colorize(result.warnings)
                     else:
+                        if sys.stdout.isatty():
                             tty.msg('No warnings.')
 
                 output.seek(0)
@@ -437,8 +475,9 @@ def solve(specs, dump=None, models=1):
 
                 # dump the raw output of the solver
                 if 'output' in dump:
+                    if sys.stdout.isatty():
                         tty.msg('Clingo output:')
-                    sys.stdout.write(result.output)
+                    colorize(result.output)
 
                 output.seek(0)
                 parser.parse(output, result)