1 files changed, 89 insertions, 16 deletions

diff --git a/tools/u_boot_pylib/terminal.py b/tools/u_boot_pylib/terminal.py
index 2cd5a54ab52..69c183e85e5 100644
--- a/tools/u_boot_pylib/terminal.py
+++ b/tools/u_boot_pylib/terminal.py
@@ -7,9 +7,12 @@
 This module handles terminal interaction including ANSI color codes.
 """
 
+from contextlib import contextmanager
+from io import StringIO
 import os
 import re
 import shutil
+import subprocess
 import sys
 
 # Selection of when we want our output to be colored
@@ -26,6 +29,13 @@ last_print_len = None
 # stackoverflow.com/questions/14693701/how-can-i-remove-the-ansi-escape-sequences-from-a-string-in-python
 ansi_escape = re.compile(r'\x1b(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
 
+# True if we are capturing console output
+CAPTURING = False
+
+# Set this to False to disable output-capturing globally
+USE_CAPTURE = True
+
+
 class PrintLine:
     """A line of text output
 
@@ -130,7 +140,8 @@ def trim_ascii_len(text, size):
     return out
 
 
-def tprint(text='', newline=True, colour=None, limit_to_line=False, bright=True):
+def tprint(text='', newline=True, colour=None, limit_to_line=False,
+           bright=True, back=None, col=None):
     """Handle a line of output to the terminal.
 
     In test mode this is recorded in a list. Otherwise it is output to the
@@ -146,9 +157,10 @@ def tprint(text='', newline=True, colour=None, limit_to_line=False, bright=True)
     if print_test_mode:
         print_test_list.append(PrintLine(text, colour, newline, bright))
     else:
-        if colour:
-            col = Color()
-            text = col.build(colour, text, bright=bright)
+        if colour is not None:
+            if not col:
+                col = Color()
+            text = col.build(colour, text, bright=bright, back=back)
         if newline:
             print(text)
             last_print_len = None
@@ -200,14 +212,23 @@ def echo_print_test_lines():
         if line.newline:
             print()
 
+def have_terminal():
+    """Check if we have an interactive terminal or not
+
+    Returns:
+        bool: true if an interactive terminal is attached
+    """
+    return os.isatty(sys.stdout.fileno())
+
 
-class Color(object):
+class Color():
     """Conditionally wraps text in ANSI color escape sequences."""
     BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
     BOLD = -1
-    BRIGHT_START = '\033[1;%dm'
-    NORMAL_START = '\033[22;%dm'
+    BRIGHT_START = '\033[1;%d%sm'
+    NORMAL_START = '\033[22;%d%sm'
     BOLD_START = '\033[1m'
+    BACK_EXTRA = ';%d'
     RESET = '\033[0m'
 
     def __init__(self, colored=COLOR_IF_TERMINAL):
@@ -224,7 +245,14 @@ class Color(object):
         except:
             self._enabled = False
 
-    def start(self, color, bright=True):
+    def enabled(self):
+        """Check if colour is enabled
+
+        Return: True if enabled, else False
+        """
+        return self._enabled
+
+    def start(self, color, bright=True, back=None):
         """Returns a start color code.
 
         Args:
@@ -235,8 +263,11 @@ class Color(object):
           color, otherwise returns empty string
         """
         if self._enabled:
+            if color == self.BOLD:
+                return self.BOLD_START
             base = self.BRIGHT_START if bright else self.NORMAL_START
-            return base % (color + 30)
+            extra = self.BACK_EXTRA % (back + 40) if back else ''
+            return base % (color + 30, extra)
         return ''
 
     def stop(self):
@@ -250,7 +281,7 @@ class Color(object):
             return self.RESET
         return ''
 
-    def build(self, color, text, bright=True):
+    def build(self, color, text, bright=True, back=None):
         """Returns text with conditionally added color escape sequences.
 
         Keyword arguments:
@@ -265,9 +296,51 @@ class Color(object):
         """
         if not self._enabled:
             return text
-        if color == self.BOLD:
-            start = self.BOLD_START
-        else:
-            base = self.BRIGHT_START if bright else self.NORMAL_START
-            start = base % (color + 30)
-        return start + text + self.RESET
+        return self.start(color, bright, back) + text + self.RESET
+
+
+# Use this to suppress stdout/stderr output:
+# with terminal.capture() as (stdout, stderr)
+#   ...do something...
+@contextmanager
+def capture():
+    global CAPTURING
+
+    capture_out, capture_err = StringIO(), StringIO()
+    old_out, old_err = sys.stdout, sys.stderr
+    try:
+        CAPTURING = True
+        sys.stdout, sys.stderr = capture_out, capture_err
+        yield capture_out, capture_err
+    finally:
+        sys.stdout, sys.stderr = old_out, old_err
+        CAPTURING = False
+        if not USE_CAPTURE:
+            sys.stdout.write(capture_out.getvalue())
+            sys.stderr.write(capture_err.getvalue())
+
+
+@contextmanager
+def pager():
+    """Simple pager for outputting lots of text
+
+    Usage:
+        with terminal.pager():
+            print(...)
+    """
+    proc = None
+    old_stdout = None
+    try:
+        less = os.getenv('PAGER')
+        if not CAPTURING and less != 'none' and have_terminal():
+            if not less:
+                less = 'less -R --quit-if-one-screen'
+            proc = subprocess.Popen(less, stdin=subprocess.PIPE, text=True,
+                                    shell=True)
+            old_stdout = sys.stdout
+            sys.stdout = proc.stdin
+        yield
+    finally:
+        if proc:
+            sys.stdout = old_stdout
+            proc.communicate()