Make nvcmp output colorful

Also refactored nicelogger to reuse the logic there.

nvchecker/lib/nicelogger.py

@@ -11,10 +11,10 @@ import sys
 import time
 import logging
 
-class TornadoLogFormatter(logging.Formatter):
-  def __init__(self, color, *args, **kwargs):
-    super().__init__(*args, **kwargs)
-    self._color = color
+class Colors:
+  def __init__(self, color=None):
+    if color is None:
+      color = support_color()
     if color:
       import curses
       curses.setupterm()
@@ -23,19 +23,32 @@ class TornadoLogFormatter(logging.Formatter):
                    curses.tigetstr("setf") or "", "ascii")
       else:
         fg_color = curses.tigetstr("setaf") or curses.tigetstr("setf") or b""
+
+      self.blue = str(curses.tparm(fg_color, 4), "ascii")
+      self.yellow = str(curses.tparm(fg_color, 3), "ascii")
+      self.green = str(curses.tparm(fg_color, 2), "ascii")
+      self.red = str(curses.tparm(fg_color, 1), "ascii")
+      self.bright_red = str(curses.tparm(fg_color, 9), "ascii")
+      self.normal = str(curses.tigetstr("sgr0"), "ascii")
+
+    else:
+      self.blue = self.yellow = self.green = self.red = self.bright_red = self.normal = ""
+
+
+class TornadoLogFormatter(logging.Formatter):
+  def __init__(self, color, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self._color = color
+    if color:
+      colors = Colors(color=color)
       self._colors = {
-        logging.DEBUG: str(curses.tparm(fg_color, 4), # Blue
-                     "ascii"),
-        logging.INFO: str(curses.tparm(fg_color, 2), # Green
-                    "ascii"),
-        logging.WARNING: str(curses.tparm(fg_color, 3), # Yellow
-                     "ascii"),
-        logging.ERROR: str(curses.tparm(fg_color, 1), # Red
-                     "ascii"),
-        logging.CRITICAL: str(curses.tparm(fg_color, 9), # Bright Red
-                     "ascii"),
+        logging.DEBUG: colors.blue,
+        logging.INFO: colors.green,
+        logging.WARNING: colors.yellow,
+        logging.ERROR: colors.red,
+        logging.CRITICAL: colors.bright_red,
       }
-      self._normal = str(curses.tigetstr("sgr0"), "ascii")
+      self._normal = colors.normal
 
   def format(self, record):
     try:
@@ -67,6 +80,18 @@ class TornadoLogFormatter(logging.Formatter):
       formatted = formatted.rstrip() + "\n" + record.exc_text
     return formatted.replace("\n", "\n    ")
 
+def support_color(stream=sys.stderr):
+  if stream.isatty():
+    try:
+      import curses
+      curses.setupterm()
+      if curses.tigetnum("colors") > 0:
+        return True
+    except:
+      import traceback
+      traceback.print_exc()
+  return False
+
 def enable_pretty_logging(level=logging.DEBUG, handler=None, color=None):
   '''
   handler: specify a handler instead of default StreamHandler
@@ -78,17 +103,8 @@ def enable_pretty_logging(level=logging.DEBUG, handler=None, color=None):
     h = logging.StreamHandler()
   else:
     h = handler
-  if color is None:
-    color = False
-    if handler is None and sys.stderr.isatty():
-      try:
-        import curses
-        curses.setupterm()
-        if curses.tigetnum("colors") > 0:
-          color = True
-      except:
-        import traceback
-        traceback.print_exc()
+  if color is None and handler is None:
+    color = support_color()
   formatter = TornadoLogFormatter(color=color)
   h.setLevel(level)
   h.setFormatter(formatter)

nvchecker/tools.py

@@ -89,4 +89,6 @@ def cmp() -> None:
       if args.quiet:
         print(name)
       else:
-        print('%s %s -> %s' % (name, oldver, newver))
+        from .lib.nicelogger import Colors, support_color
+        c = Colors(support_color(sys.stdout))
+        print(f'{name} {c.red}{oldver}{c.normal} -> {c.green}{newver}{c.normal}')