Merge pull request #850 from python-cmd2/truncate_string

Truncate line

1 files changed, 71 insertions, 14 deletions

diff --git a/cmd2/utils.py b/cmd2/utils.py
index ffbe5a64..3e154bd2 100644
--- a/cmd2/utils.py
+++ b/cmd2/utils.py
@@ -638,7 +638,7 @@ class TextAlignment(Enum):
 
 
 def align_text(text: str, alignment: TextAlignment, *, fill_char: str = ' ',
-               width: Optional[int] = None, tab_width: int = 4) -> str:
+               width: Optional[int] = None, tab_width: int = 4, truncate: bool = False) -> str:
     """
     Align text for display within a given width. Supports characters with display widths greater than 1.
     ANSI style sequences are safely ignored and do not count toward the display width. This means colored text is
@@ -652,15 +652,24 @@ def align_text(text: str, alignment: TextAlignment, *, fill_char: str = ' ',
     :param width: display width of the aligned text. Defaults to width of the terminal.
     :param tab_width: any tabs in the text will be replaced with this many spaces. if fill_char is a tab, then it will
                       be converted to a space.
+    :param truncate: if True, then each line will be shortened to fit within the display width. The truncated
+                     portions are replaced by a '…' character. Defaults to False.
     :return: aligned text
     :raises: TypeError if fill_char is more than one character
              ValueError if text or fill_char contains an unprintable character
+             ValueError if width is less than 1
     """
     import io
     import shutil
 
     from . import ansi
 
+    if width is None:
+        width = shutil.get_terminal_size().columns
+
+    if width < 1:
+        raise ValueError("width must be at least 1")
+
     # Handle tabs
     text = text.replace('\t', ' ' * tab_width)
     if fill_char == '\t':
@@ -678,23 +687,21 @@ def align_text(text: str, alignment: TextAlignment, *, fill_char: str = ' ',
     else:
         lines = ['']
 
-    if width is None:
-        width = shutil.get_terminal_size().columns
-
     text_buf = io.StringIO()
 
     for index, line in enumerate(lines):
         if index > 0:
             text_buf.write('\n')
 
-        # Use style_aware_wcswidth to support characters with display widths
-        # greater than 1 as well as ANSI style sequences
+        if truncate:
+            line = truncate_line(line, width)
+
         line_width = ansi.style_aware_wcswidth(line)
         if line_width == -1:
             raise(ValueError("Text to align contains an unprintable character"))
 
-        # Check if line is wider than the desired final width
-        if width <= line_width:
+        elif line_width >= width:
+            # No need to add fill characters
             text_buf.write(line)
             continue
 
@@ -725,7 +732,8 @@ def align_text(text: str, alignment: TextAlignment, *, fill_char: str = ' ',
     return text_buf.getvalue()
 
 
-def align_left(text: str, *, fill_char: str = ' ', width: Optional[int] = None, tab_width: int = 4) -> str:
+def align_left(text: str, *, fill_char: str = ' ', width: Optional[int] = None,
+               tab_width: int = 4, truncate: bool = False) -> str:
     """
     Left align text for display within a given width. Supports characters with display widths greater than 1.
     ANSI style sequences are safely ignored and do not count toward the display width. This means colored text is
@@ -736,14 +744,19 @@ def align_left(text: str, *, fill_char: str = ' ', width: Optional[int] = None,
     :param width: display width of the aligned text. Defaults to width of the terminal.
     :param tab_width: any tabs in the text will be replaced with this many spaces. if fill_char is a tab, then it will
                       be converted to a space.
+    :param truncate: if True, then text will be shortened to fit within the display width. The truncated portion is
+                     replaced by a '…' character. Defaults to False.
     :return: left-aligned text
     :raises: TypeError if fill_char is more than one character
              ValueError if text or fill_char contains an unprintable character
+             ValueError if width is less than 1
     """
-    return align_text(text, TextAlignment.LEFT, fill_char=fill_char, width=width, tab_width=tab_width)
+    return align_text(text, TextAlignment.LEFT, fill_char=fill_char, width=width,
+                      tab_width=tab_width, truncate=truncate)
 
 
-def align_center(text: str, *, fill_char: str = ' ', width: Optional[int] = None, tab_width: int = 4) -> str:
+def align_center(text: str, *, fill_char: str = ' ', width: Optional[int] = None,
+                 tab_width: int = 4, truncate: bool = False) -> str:
     """
     Center text for display within a given width. Supports characters with display widths greater than 1.
     ANSI style sequences are safely ignored and do not count toward the display width. This means colored text is
@@ -754,14 +767,19 @@ def align_center(text: str, *, fill_char: str = ' ', width: Optional[int] = None
     :param width: display width of the aligned text. Defaults to width of the terminal.
     :param tab_width: any tabs in the text will be replaced with this many spaces. if fill_char is a tab, then it will
                       be converted to a space.
+    :param truncate: if True, then text will be shortened to fit within the display width. The truncated portion is
+                     replaced by a '…' character. Defaults to False.
     :return: centered text
     :raises: TypeError if fill_char is more than one character
              ValueError if text or fill_char contains an unprintable character
+             ValueError if width is less than 1
     """
-    return align_text(text, TextAlignment.CENTER, fill_char=fill_char, width=width, tab_width=tab_width)
+    return align_text(text, TextAlignment.CENTER, fill_char=fill_char, width=width,
+                      tab_width=tab_width, truncate=truncate)
 
 
-def align_right(text: str, *, fill_char: str = ' ', width: Optional[int] = None, tab_width: int = 4) -> str:
+def align_right(text: str, *, fill_char: str = ' ', width: Optional[int] = None,
+                tab_width: int = 4, truncate: bool = False) -> str:
     """
     Right align text for display within a given width. Supports characters with display widths greater than 1.
     ANSI style sequences are safely ignored and do not count toward the display width. This means colored text is
@@ -772,8 +790,47 @@ def align_right(text: str, *, fill_char: str = ' ', width: Optional[int] = None,
     :param width: display width of the aligned text. Defaults to width of the terminal.
     :param tab_width: any tabs in the text will be replaced with this many spaces. if fill_char is a tab, then it will
                       be converted to a space.
+    :param truncate: if True, then text will be shortened to fit within the display width. The truncated portion is
+                     replaced by a '…' character. Defaults to False.
     :return: right-aligned text
     :raises: TypeError if fill_char is more than one character
              ValueError if text or fill_char contains an unprintable character
+             ValueError if width is less than 1
+    """
+    return align_text(text, TextAlignment.RIGHT, fill_char=fill_char, width=width,
+                      tab_width=tab_width, truncate=truncate)
+
+
+def truncate_line(line: str, max_width: int, *, tab_width: int = 4) -> str:
     """
-    return align_text(text, TextAlignment.RIGHT, fill_char=fill_char, width=width, tab_width=tab_width)
+    Truncate a single line to fit within a given display width. Any portion of the string that is truncated
+    is replaced by a '…' character. Supports characters with display widths greater than 1. ANSI style sequences are
+    safely ignored and do not count toward the display width. This means colored text is supported.
+
+    :param line: text to truncate
+    :param max_width: the maximum display width the resulting string is allowed to have
+    :param tab_width: any tabs in the text will be replaced with this many spaces
+    :return: line that has a display width less than or equal to width
+    :raises: ValueError if text contains an unprintable character like a new line
+             ValueError if max_width is less than 1
+    """
+    from . import ansi
+
+    # Handle tabs
+    line = line.replace('\t', ' ' * tab_width)
+
+    if ansi.style_aware_wcswidth(line) == -1:
+        raise (ValueError("text contains an unprintable character"))
+
+    if max_width < 1:
+        raise ValueError("max_width must be at least 1")
+
+    if ansi.style_aware_wcswidth(line) > max_width:
+        # Remove characters until we fit. Leave room for the ellipsis.
+        line = line[:max_width - 1]
+        while ansi.style_aware_wcswidth(line) > max_width - 1:
+            line = line[:-1]
+
+        line += "\N{HORIZONTAL ELLIPSIS}"
+
+    return line