Add support for colorization in loopy terminal output

loopy/kernel/__init__.py

@@ -1012,6 +1012,9 @@ class LoopKernel(RecordWithoutPickling):
 
         printed_insn_ids = set()
 
+        Fore = self.options._fore
+        Style = self.options._style
+
         def print_insn(insn):
             if insn.id in printed_insn_ids:
                 return
@@ -1035,7 +1038,7 @@ class LoopKernel(RecordWithoutPickling):
 
             loop_list = ",".join(sorted(kernel.insn_inames(insn)))
 
-            options = [insn.id]
+            options = [Fore.GREEN+insn.id+Style.RESET_ALL]
             if insn.priority:
                 options.append("priority=%d" % insn.priority)
             if insn.tags:
@@ -1048,12 +1051,15 @@ class LoopKernel(RecordWithoutPickling):
             if len(loop_list) > loop_list_width:
                 lines.append("[%s]" % loop_list)
                 lines.append("%s%s <- %s   # %s" % (
-                    (loop_list_width+2)*" ", lhs,
-                    rhs, ", ".join(options)))
+                    (loop_list_width+2)*" ", Fore.BLUE+lhs+Style.RESET_ALL,
+                    Fore.MAGENTA+rhs+Style.RESET_ALL,
+                    ", ".join(options)))
             else:
                 lines.append("[%s]%s%s <- %s   # %s" % (
                     loop_list, " "*(loop_list_width-len(loop_list)),
-                    lhs, rhs, ",".join(options)))
+                    Fore.BLUE + lhs + Style.RESET_ALL,
+                    Fore.MAGENTA+rhs+Style.RESET_ALL,
+                    ",".join(options)))
 
             lines.extend(trailing)
 

loopy/options.py

@@ -27,6 +27,11 @@ from pytools import Record
 import re
 
 
+class _ColoramaStub(object):
+    def getattribute(self, name):
+        return ""
+
+
 class Options(Record):
     """
     Unless otherwise specified, these options are Boolean-valued
@@ -105,6 +110,11 @@ class Options(Record):
 
         Options to pass to the OpenCL compiler when building the kernel.
         A list of strings.
+
+    .. attribute:: allow_terminal_colors
+
+        A :class:`bool`. Whether to allow colors in terminal output
+
     """
 
     def __init__(
@@ -124,6 +134,7 @@ class Options(Record):
             write_wrapper=False, highlight_wrapper=False,
             write_cl=False, highlight_cl=False,
             edit_cl=False, cl_build_options=[],
+            allow_terminal_colors=True
             ):
         Record.__init__(
                 self,
@@ -137,6 +148,7 @@ class Options(Record):
                 write_wrapper=write_wrapper, highlight_wrapper=highlight_wrapper,
                 write_cl=write_cl, highlight_cl=highlight_cl,
                 edit_cl=edit_cl, cl_build_options=cl_build_options,
+                allow_terminal_colors=allow_terminal_colors,
                 )
 
     def update(self, other):
@@ -150,6 +162,30 @@ class Options(Record):
         for field_name in sorted(self.__class__.fields):
             key_builder.rec(key_hash, getattr(self, field_name))
 
+    @property
+    def _fore(self):
+        if self.allow_terminal_colors:
+            import colorama
+            return colorama.Fore
+        else:
+            return _ColoramaStub()
+
+    @property
+    def _back(self):
+        if self.allow_terminal_colors:
+            import colorama
+            return colorama.Back
+        else:
+            return _ColoramaStub()
+
+    @property
+    def _style(self):
+        if self.allow_terminal_colors:
+            import colorama
+            return colorama.Style
+        else:
+            return _ColoramaStub()
+
 
 KEY_VAL_RE = re.compile("^([a-zA-Z0-9]+)=(.*)$")
 

setup.py

@@ -42,6 +42,7 @@ setup(name="loo.py",
           "cgen>=2016.1",
           "islpy>=2016.1.2",
           "six>=1.8.0",
+          "colorama",
           ],
 
       extras_require={