Unify handling of console color output. NFC

Followup to #25502.

Author: sbc100

tools/cmdline.py

@@ -510,9 +510,9 @@ def consume_arg_file():
     elif check_flag('--threadprofiler'):
       settings_changes.append('PTHREADS_PROFILING=1')
     elif arg in ('-fcolor-diagnostics', '-fdiagnostics-color', '-fdiagnostics-color=always'):
-      diagnostics.color_enabled = True
-    elif arg in ('-fno-color-diagnostics', '-fdiagnostics-color=never'):
-      diagnostics.color_enabled = False
+      colored_logger.enable()
+    elif arg in ('-fno-color-diagnostics', '-fno-diagnostics-color', '-fdiagnostics-color=never'):
+      colored_logger.disable()
     elif arg == '-fno-exceptions':
       settings.DISABLE_EXCEPTION_CATCHING = 1
       settings.DISABLE_EXCEPTION_THROWING = 1
@@ -568,9 +568,6 @@ def consume_arg_file():
       settings.USE_PTHREADS = 0
     elif arg == '-pthreads':
       exit_with_error('unrecognized command-line option `-pthreads`; did you mean `-pthread`?')
-    elif arg in ('-fno-diagnostics-color', '-fdiagnostics-color=never'):
-      colored_logger.disable()
-      diagnostics.color_enabled = False
     elif arg == '-fno-rtti':
       settings.USE_RTTI = 0
     elif arg == '-frtti':

tools/colored_logger.py

@@ -4,11 +4,59 @@
 # found in the LICENSE file.
 
 """Enables colored logger just by importing this module
+
+Also, provides utiliy functions to use ANSI colors in the terminal.
 """
 
 import ctypes
 import sys
 import logging
+from functools import wraps
+
+
+# ANSI colors
+RED = 1
+GREEN = 2
+YELLOW = 3
+BLUE = 4
+MAGENTA = 5
+CYAN = 6
+WHITE = 7
+
+color_enabled = False
+
+
+def output_color(color):
+  assert color_enabled
+  return '\033[3%sm' % color
+
+
+def bold():
+  assert color_enabled
+  return '\033[1m'
+
+
+def reset_color():
+  assert color_enabled
+  return '\033[0m'
+
+
+def with_bold_color(color, string):
+  if not color_enabled:
+    return string
+  return output_color(color) + bold() + string + reset_color()
+
+
+def with_color(color, string):
+  if not color_enabled:
+    return string
+  return output_color(color) + string + reset_color()
+
+
+def with_bold(string):
+  if not color_enabled:
+    return string
+  return bold() + string + reset_color()
 
 
 def ansi_color_available():
@@ -34,30 +82,35 @@ def ansi_color_available():
 
 def add_coloring_to_emit_ansi(fn):
   # add methods we need to the class
+  @wraps(fn)
   def new(*args):
     levelno = args[1].levelno
+    color = None
     if levelno >= 40:
-      color = '\x1b[31m' # red
+      color = RED
     elif levelno >= 30:
-      color = '\x1b[33m' # yellow
+      color = YELLOW
     elif levelno >= 20:
-      color = '\x1b[32m' # green
+      color = GREEN
     elif levelno >= 10:
-      color = '\x1b[35m' # pink
-    else:
-      color = '\x1b[0m' # normal
-    args[1].msg = color + args[1].msg + '\x1b[0m'  # normal
+      color = MAGENTA
+    if color:
+      args[1].msg = with_color(color, args[1].msg)
     return fn(*args)
 
   new.orig_func = fn
   return new
 
 
-def enable():
-  if ansi_color_available():
+def enable(force=False):
+  global color_enabled
+  if force or ansi_color_available():
     logging.StreamHandler.emit = add_coloring_to_emit_ansi(logging.StreamHandler.emit)
+    color_enabled = True
 
 
 def disable():
+  global color_enabled
   if hasattr(logging.StreamHandler.emit, 'orig_func'):
     logging.StreamHandler.emit = logging.StreamHandler.emit.orig_func
+  color_enabled = False

tools/diagnostics.py

@@ -13,27 +13,17 @@
 
 from . import colored_logger
 
-color_enabled = colored_logger.ansi_color_available()
 logger = logging.getLogger('diagnostics')
 tool_name = os.path.splitext(os.path.basename(sys.argv[0]))[0]
 
 # diagnostic levels
 WARN = 1
 ERROR = 2
 
-# available (ANSI) colors
-RED = 1
-GREEN = 2
-YELLOW = 3
-BLUE = 4
-MAGENTA = 5
-CYAN = 6
-WHITE = 7
-
 # color for use for each diagnostic level
 level_colors = {
-    WARN: MAGENTA,
-    ERROR: RED,
+    WARN: colored_logger.MAGENTA,
+    ERROR: colored_logger.RED,
 }
 
 level_prefixes = {
@@ -42,48 +32,19 @@
 }
 
 
-def output_color(color):
-  assert color_enabled
-  return '\033[3%sm' % color
-
-
-def bold():
-  assert color_enabled
-  return '\033[1m'
-
-
-def reset_color():
-  assert color_enabled
-  return '\033[0m'
-
-
 def diag(level, msg, *args):
   # Format output message as:
   # <tool>: <level>: msg
-  # With the `<level>:` part being colored accordingly.
-  sys.stderr.write(tool_name + ': ')
-
-  if color_enabled:
-    output = output_color(level_colors[level]) + bold()
-    if output:
-      sys.stderr.write(output)
-
-  sys.stderr.write(level_prefixes[level])
-
-  if color_enabled:
-    output = reset_color() + bold()
-    if output:
-      sys.stderr.write(output)
-
+  # With the `<level>:` part being colored accordingly, and the message itself in bold.
+  prefix = level_prefixes[level]
+  color = level_colors[level]
   if args:
     msg = msg % args
-  sys.stderr.write(str(msg))
-  sys.stderr.write('\n')
 
-  if color_enabled:
-    output = reset_color()
-    if output:
-      sys.stderr.write(output)
+  # Add colors
+  prefix = colored_logger.with_bold_color(color, prefix)
+  msg = colored_logger.with_bold(msg)
+  sys.stderr.write(f'{tool_name}: {prefix}{msg}\n')
 
 
 def error(msg, *args):