From 8bc1092f41741a3ef1de9d360cd98ccb3b317962 Mon Sep 17 00:00:00 2001
From: Todd Gamblin <tgamblin@llnl.gov>
Date: Sat, 10 Aug 2019 14:27:55 -0700
Subject: [PATCH] concretizer: colorize ASP output

---
 lib/spack/spack/solver/asp.py | 53 ++++++++++++++++++++++++++++++-----
 1 file changed, 46 insertions(+), 7 deletions(-)

diff --git a/lib/spack/spack/solver/asp.py b/lib/spack/spack/solver/asp.py
index 30268e5224..55bf38ec50 100644
--- a/lib/spack/spack/solver/asp.py
+++ b/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:
-            tty.msg('ASP program:')
-            sys.stdout.write(result.asp)
+            if sys.stdout.isatty():
+                tty.msg('ASP program:')
+            colorize(result.asp)
 
         with tempfile.TemporaryFile("w+") as output:
             with tempfile.TemporaryFile() as warnings:
@@ -427,18 +463,21 @@ def solve(specs, dump=None, models=1):
                 # dump any warnings generated by the solver
                 if 'warnings' in dump:
                     if result.warnings:
-                        tty.msg('Clingo gave the following warnings:')
-                        sys.stdout.write(result.warnings)
+                        if sys.stdout.isatty():
+                            tty.msg('Clingo gave the following warnings:')
+                        colorize(result.warnings)
                     else:
-                        tty.msg('No warnings.')
+                        if sys.stdout.isatty():
+                            tty.msg('No warnings.')
 
                 output.seek(0)
                 result.output = output.read()
 
                 # dump the raw output of the solver
                 if 'output' in dump:
-                    tty.msg('Clingo output:')
-                    sys.stdout.write(result.output)
+                    if sys.stdout.isatty():
+                        tty.msg('Clingo output:')
+                    colorize(result.output)
 
                 output.seek(0)
                 parser.parse(output, result)