patman: Drop references to __future__

[oweals/u-boot.git] / tools / patman / terminal.py

diff --git a/tools/patman/terminal.py b/tools/patman/terminal.py
index e78a7c14f5bfd87f9d7344d9c3ada90d1df973b5..c709438bdc4f737b8509e6e69f6565381785994e 100644 (file)
--- a/tools/patman/terminal.py
+++ b/tools/patman/terminal.py
@@ -1,7 +1,6 @@
+# SPDX-License-Identifier: GPL-2.0+
 # Copyright (c) 2011 The Chromium OS Authors.
 #
-# SPDX-License-Identifier:     GPL-2.0+
-#
 
 """Terminal utilities
 
@@ -9,6 +8,8 @@ This module handles terminal interaction including ANSI color codes.
 """
 
 import os
+import re
+import shutil
 import sys
 
 # Selection of when we want our output to be colored
@@ -18,6 +19,13 @@ COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3)
 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-?]*[ -/]*[@-~])')
+
 class PrintLine:
     """A line of text output
 
@@ -35,7 +43,86 @@ class PrintLine:
         return 'newline=%s, colour=%s, text=%s' % (self.newline, self.colour,
                 self.text)
 
-def Print(text='', newline=True, colour=None):
+def CalcAsciiLen(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.Color(Color.RED, 'abc')
+    >>> len(text)
+    14
+    >>> CalcAsciiLen(text)
+    3
+    >>>
+    >>> text += 'def'
+    >>> CalcAsciiLen(text)
+    6
+    >>> text += col.Color(Color.RED, 'abc')
+    >>> CalcAsciiLen(text)
+    9
+    """
+    result = ansi_escape.sub('', text)
+    return len(result)
+
+def TrimAsciiLen(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.Color(Color.RED, 'abc')
+    >>> len(text)
+    14
+    >>> CalcAsciiLen(TrimAsciiLen(text, 4))
+    3
+    >>> CalcAsciiLen(TrimAsciiLen(text, 2))
+    2
+    >>> text += 'def'
+    >>> CalcAsciiLen(TrimAsciiLen(text, 4))
+    4
+    >>> text += col.Color(Color.RED, 'ghi')
+    >>> CalcAsciiLen(TrimAsciiLen(text, 7))
+    7
+    """
+    if CalcAsciiLen(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 Print(text='', newline=True, colour=None, limit_to_line=False):
     """Handle a line of output to the terminal.
 
     In test mode this is recorded in a list. Otherwise it is output to the
@@ -46,15 +133,31 @@ def Print(text='', newline=True, colour=None):
         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, newline, colour))
     else:
         if colour:
             col = Color()
             text = col.Color(colour, text)
-        print text,
         if newline:
-            print
+            print(text)
+            last_print_len = None
+        else:
+            if limit_to_line:
+                cols = shutil.get_terminal_size().columns
+                text = TrimAsciiLen(text, cols)
+            print(text, end='', flush=True)
+            last_print_len = CalcAsciiLen(text)
+
+def PrintClear():
+    """Clear a previously line that was printed with no newline"""
+    global last_print_len
+
+    if last_print_len:
+        print('\r%s\r' % (' '* last_print_len), end='', flush=True)
+        last_print_len = None
 
 def SetPrintTestMode():
     """Go into test mode, where all printing is recorded"""
@@ -79,11 +182,11 @@ def EchoPrintTestLines():
     for line in print_test_list:
         if line.colour:
             col = Color()
-            print col.Color(line.colour, line.text),
+            print(col.Color(line.colour, line.text), end='')
         else:
-            print line.text,
+            print(line.text, end='')
         if line.newline:
-            print
+            print()
 
 
 class Color(object):
@@ -125,7 +228,7 @@ class Color(object):
         return ''
 
     def Stop(self):
-        """Retruns a stop color code.
+        """Returns a stop color code.
 
         Returns:
           If color is enabled, returns an ANSI color reset sequence,