1 files changed, 346 insertions, 0 deletions
diff --git a/tools/u_boot_pylib/terminal.py b/tools/u_boot_pylib/terminal.py
new file mode 100644
index 00000000000..69c183e85e5
--- /dev/null
+++ b/tools/u_boot_pylib/terminal.py
@@ -0,0 +1,346 @@
+# SPDX-License-Identifier: GPL-2.0+
+# Copyright (c) 2011 The Chromium OS Authors.
+#
+
+"""Terminal utilities
+
+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
+COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3)
+
+# Initially, we are set up to print to the terminal
+print_test_mode = False
+print_test_list = []
+
+# The length of the last line printed without a newline
+last_print_len = None
+
+# credit:
+# 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
+
+    Members:
+        text: Text line that was printed
+        newline: True to output a newline after the text
+        colour: Text colour to use
+    """
+    def __init__(self, text, colour, newline=True, bright=True):
+        self.text = text
+        self.newline = newline
+        self.colour = colour
+        self.bright = bright
+
+    def __eq__(self, other):
+        return (self.text == other.text and
+                self.newline == other.newline and
+                self.colour == other.colour and
+                self.bright == other.bright)
+
+    def __str__(self):
+        return ("newline=%s, colour=%s, bright=%d, text='%s'" %
+                (self.newline, self.colour, self.bright, self.text))
+
+
+def calc_ascii_len(text):
+    """Calculate the length of a string, ignoring any ANSI sequences
+
+    When displayed on a terminal, ANSI sequences don't take any space, so we
+    need to ignore them when calculating the length of a string.
+
+    Args:
+        text: Text to check
+
+    Returns:
+        Length of text, after skipping ANSI sequences
+
+    >>> col = Color(COLOR_ALWAYS)
+    >>> text = col.build(Color.RED, 'abc')
+    >>> len(text)
+    14
+    >>> calc_ascii_len(text)
+    3
+    >>>
+    >>> text += 'def'
+    >>> calc_ascii_len(text)
+    6
+    >>> text += col.build(Color.RED, 'abc')
+    >>> calc_ascii_len(text)
+    9
+    """
+    result = ansi_escape.sub('', text)
+    return len(result)
+
+def trim_ascii_len(text, size):
+    """Trim a string containing ANSI sequences to the given ASCII length
+
+    The string is trimmed with ANSI sequences being ignored for the length
+    calculation.
+
+    >>> col = Color(COLOR_ALWAYS)
+    >>> text = col.build(Color.RED, 'abc')
+    >>> len(text)
+    14
+    >>> calc_ascii_len(trim_ascii_len(text, 4))
+    3
+    >>> calc_ascii_len(trim_ascii_len(text, 2))
+    2
+    >>> text += 'def'
+    >>> calc_ascii_len(trim_ascii_len(text, 4))
+    4
+    >>> text += col.build(Color.RED, 'ghi')
+    >>> calc_ascii_len(trim_ascii_len(text, 7))
+    7
+    """
+    if calc_ascii_len(text) < size:
+        return text
+    pos = 0
+    out = ''
+    left = size
+
+    # Work through each ANSI sequence in turn
+    for m in ansi_escape.finditer(text):
+        # Find the text before the sequence and add it to our string, making
+        # sure it doesn't overflow
+        before = text[pos:m.start()]
+        toadd = before[:left]
+        out += toadd
+
+        # Figure out how much non-ANSI space we have left
+        left -= len(toadd)
+
+        # Add the ANSI sequence and move to the position immediately after it
+        out += m.group()
+        pos = m.start() + len(m.group())
+
+    # Deal with text after the last ANSI sequence
+    after = text[pos:]
+    toadd = after[:left]
+    out += toadd
+
+    return out
+
+
+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
+    terminal.
+
+    Args:
+        text: Text to print
+        newline: True to add a new line at the end of the text
+        colour: Colour to use for the text
+    """
+    global last_print_len
+
+    if print_test_mode:
+        print_test_list.append(PrintLine(text, colour, newline, bright))
+    else:
+        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
+        else:
+            if limit_to_line:
+                cols = shutil.get_terminal_size().columns
+                text = trim_ascii_len(text, cols)
+            print(text, end='', flush=True)
+            last_print_len = calc_ascii_len(text)
+
+def print_clear():
+    """Clear a previously line that was printed with no newline"""
+    global last_print_len
+
+    if last_print_len:
+        if print_test_mode:
+            print_test_list.append(PrintLine(None, None, None, None))
+        else:
+            print('\r%s\r' % (' '* last_print_len), end='', flush=True)
+            last_print_len = None
+
+def set_print_test_mode(enable=True):
+    """Go into test mode, where all printing is recorded"""
+    global print_test_mode
+
+    print_test_mode = enable
+    get_print_test_lines()
+
+def get_print_test_lines():
+    """Get a list of all lines output through tprint()
+
+    Returns:
+        A list of PrintLine objects
+    """
+    global print_test_list
+
+    ret = print_test_list
+    print_test_list = []
+    return ret
+
+def echo_print_test_lines():
+    """Print out the text lines collected"""
+    for line in print_test_list:
+        if line.colour:
+            col = Color()
+            print(col.build(line.colour, line.text), end='')
+        else:
+            print(line.text, end='')
+        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():
+    """Conditionally wraps text in ANSI color escape sequences."""
+    BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
+    BOLD = -1
+    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):
+        """Create a new Color object, optionally disabling color output.
+
+        Args:
+          enabled: True if color output should be enabled. If False then this
+            class will not add color codes at all.
+        """
+        try:
+            self._enabled = (colored == COLOR_ALWAYS or
+                    (colored == COLOR_IF_TERMINAL and
+                     os.isatty(sys.stdout.fileno())))
+        except:
+            self._enabled = False
+
+    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:
+          color: Color to use, .e.g BLACK, RED, etc.
+
+        Returns:
+          If color is enabled, returns an ANSI sequence to start the given
+          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
+            extra = self.BACK_EXTRA % (back + 40) if back else ''
+            return base % (color + 30, extra)
+        return ''
+
+    def stop(self):
+        """Returns a stop color code.
+
+        Returns:
+          If color is enabled, returns an ANSI color reset sequence,
+          otherwise returns empty string
+        """
+        if self._enabled:
+            return self.RESET
+        return ''
+
+    def build(self, color, text, bright=True, back=None):
+        """Returns text with conditionally added color escape sequences.
+
+        Keyword arguments:
+          color: Text color -- one of the color constants defined in this
+                  class.
+          text: The text to color.
+
+        Returns:
+          If self._enabled is False, returns the original text. If it's True,
+          returns text with color escape sequences based on the value of
+          color.
+        """
+        if not self._enabled:
+            return text
+        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()