pmbootstrap: add color support (MR 2090)

Can be disabled by setting the $NO_COLOR environment variable

pmb/config/__init__.py

@@ -123,6 +123,21 @@ is_interactive = sys.stdout.isatty() and \
     sys.stdin.isatty()
 
 
+# ANSI escape codes to highlight stdout
+styles = {
+    "BLUE": '\033[94m',
+    "BOLD": '\033[1m',
+    "GREEN": '\033[92m',
+    "RED": '\033[91m',
+    "YELLOW": '\033[93m',
+    "END": '\033[0m'
+}
+
+if "NO_COLOR" in os.environ:
+    for style in styles.keys():
+        styles[style] = ""
+
+
 # List of available locales taken from musl-locales package; see
 # https://pkgs.alpinelinux.org/contents?name=musl-locales
 locales = [

pmb/helpers/cli.py

@@ -50,13 +50,17 @@ def ask(args, question="Continue?", choices=["y", "n"], default="n",
     :param validation_regex: if set, keep asking until regex matches
     :param complete: set to a list to enable tab completion
     """
+    styles = pmb.config.styles
+
     while True:
         date = datetime.datetime.now().strftime("%H:%M:%S")
-        question_full = "[" + date + "] " + question
+        line = question
         if choices:
-            question_full += " (" + str.join("/", choices) + ")"
+            line += f" ({str.join('/', choices)})"
         if default:
-            question_full += " [" + str(default) + "]"
+            line += f" [{default}]"
+        line_color = f"[{date}] {styles['BOLD']}{line}{styles['END']}"
+        line = f"[{date}] {line}"
 
         if complete:
             readline.parse_and_bind('tab: complete')
@@ -67,7 +71,7 @@ def ask(args, question="Continue?", choices=["y", "n"], default="n",
             readline.set_completer(
                 ReadlineTabCompleter(complete).completer_func)
 
-        ret = input(question_full + ": ")
+        ret = input(f"{line_color}: ")
 
         # Stop completing (question is answered)
         if complete:
@@ -79,7 +83,7 @@ def ask(args, question="Continue?", choices=["y", "n"], default="n",
         if ret == "":
             ret = str(default)
 
-        args.logfd.write(question_full + " " + ret + "\n")
+        args.logfd.write(f"{line}: {ret}\n")
         args.logfd.flush()
 
         # Validate with regex

pmb/helpers/logging.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import sys
+import pmb.config
 
 
 class log_handler(logging.StreamHandler):
@@ -20,7 +21,33 @@ class log_handler(logging.StreamHandler):
                 not self._args.quiet and
                     record.levelno >= logging.INFO):
                 stream = self.stream
-                stream.write(msg)
+
+                styles = pmb.config.styles
+
+                msg_col = (
+                    msg.replace(
+                        "NOTE:",
+                        f"{styles['BLUE']}NOTE:{styles['END']}",
+                        1,
+                    )
+                    .replace(
+                        "WARNING:",
+                        f"{styles['YELLOW']}WARNING:{styles['END']}",
+                        1,
+                    )
+                    .replace(
+                        "ERROR:",
+                        f"{styles['RED']}ERROR:{styles['END']}",
+                        1,
+                    )
+                    .replace(
+                        "DONE!",
+                        f"{styles['GREEN']}DONE!{styles['END']}",
+                        1,
+                    )
+                )
+
+                stream.write(msg_col)
                 stream.write(self.terminator)
                 self.flush()