truncate_line not handling style sequences

CHANGELOG.md

@@ -2,6 +2,7 @@
 * Bug Fixes
     * Corrected issue where the actual new value was not always being printed in do_set. This occurred in cases where
       the typed value differed from what the setter had converted it to.
+    * Fixed bug where ANSI style sequences were not correctly handled in `utils.truncate_line()`. 
 * Enhancements
     * Renamed set command's `-l/--long` flag to `-v/--verbose` for consistency with help and history commands.
 

cmd2/ansi.py

@@ -118,6 +118,7 @@ class bg(ColorBase):
 INTENSITY_BRIGHT = Style.BRIGHT
 INTENSITY_DIM = Style.DIM
 INTENSITY_NORMAL = Style.NORMAL
+
 # ANSI style sequences not provided by colorama
 UNDERLINE_ENABLE = colorama.ansi.code_to_chars(4)
 UNDERLINE_DISABLE = colorama.ansi.code_to_chars(24)

cmd2/constants.py

@@ -15,6 +15,9 @@
 
 LINE_FEED = '\n'
 
+# One character ellipsis
+HORIZONTAL_ELLIPSIS = '…'
+
 DEFAULT_SHORTCUTS = {'?': 'help', '!': 'shell', '@': 'run_script', '@@': '_relative_run_script'}
 
 # Used as the command name placeholder in disabled command messages.

cmd2/utils.py

@@ -11,7 +11,7 @@
 import threading
 import unicodedata
 from enum import Enum
-from typing import Any, Callable, Iterable, List, Optional, TextIO, Union
+from typing import Any, Callable, Dict, Iterable, List, Optional, TextIO, Union
 
 from . import constants
 
@@ -682,8 +682,8 @@ def align_text(text: str, alignment: TextAlignment, *, fill_char: 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
-    supported. If text has line breaks, then each line is aligned independently.
+    ANSI style sequences do not count toward the display width. If text has line breaks, then each line is aligned
+    independently.
 
     There are convenience wrappers around this function: align_left(), align_center(), and align_right()
 
@@ -696,7 +696,7 @@ def align_text(text: str, alignment: TextAlignment, *, fill_char: str = ' ',
     :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
+    :raises: TypeError if fill_char is more than one character (not including ANSI style sequences)
              ValueError if text or fill_char contains an unprintable character
              ValueError if width is less than 1
     """
@@ -716,7 +716,7 @@ def align_text(text: str, alignment: TextAlignment, *, fill_char: str = ' ',
     if fill_char == '\t':
         fill_char = ' '
 
-    if len(fill_char) != 1:
+    if len(ansi.strip_style(fill_char)) != 1:
         raise TypeError("Fill character must be exactly one character long")
 
     fill_char_width = ansi.style_aware_wcswidth(fill_char)
@@ -777,8 +777,8 @@ 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
-    supported. If text has line breaks, then each line is aligned independently.
+    ANSI style sequences do not count toward the display width. If text has line breaks, then each line is aligned
+    independently.
 
     :param text: text to left align (can contain multiple lines)
     :param fill_char: character that fills the alignment gap. Defaults to space. (Cannot be a line breaking character)
@@ -788,7 +788,7 @@ def align_left(text: str, *, fill_char: str = ' ', width: Optional[int] = None,
     :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
+    :raises: TypeError if fill_char is more than one character (not including ANSI style sequences)
              ValueError if text or fill_char contains an unprintable character
              ValueError if width is less than 1
     """
@@ -800,8 +800,8 @@ 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
-    supported. If text has line breaks, then each line is aligned independently.
+    ANSI style sequences do not count toward the display width. If text has line breaks, then each line is aligned
+    independently.
 
     :param text: text to center (can contain multiple lines)
     :param fill_char: character that fills the alignment gap. Defaults to space. (Cannot be a line breaking character)
@@ -811,7 +811,7 @@ def align_center(text: str, *, fill_char: str = ' ', width: Optional[int] = None
     :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
+    :raises: TypeError if fill_char is more than one character (not including ANSI style sequences)
              ValueError if text or fill_char contains an unprintable character
              ValueError if width is less than 1
     """
@@ -823,8 +823,8 @@ 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
-    supported. If text has line breaks, then each line is aligned independently.
+    ANSI style sequences do not count toward the display width. If text has line breaks, then each line is aligned
+    independently.
 
     :param text: text to right align (can contain multiple lines)
     :param fill_char: character that fills the alignment gap. Defaults to space. (Cannot be a line breaking character)
@@ -834,7 +834,7 @@ def align_right(text: str, *, fill_char: str = ' ', width: Optional[int] = None,
     :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
+    :raises: TypeError if fill_char is more than one character (not including ANSI style sequences)
              ValueError if text or fill_char contains an unprintable character
              ValueError if width is less than 1
     """
@@ -845,8 +845,15 @@ def align_right(text: str, *, fill_char: str = ' ', width: Optional[int] = None,
 def truncate_line(line: str, max_width: int, *, tab_width: int = 4) -> str:
     """
     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.
+    is replaced by a '…' character. Supports characters with display widths greater than 1. ANSI style sequences
+    do not count toward the display width.
+
+    If there are ANSI style sequences in the string after where truncation occurs, this function will append them
+    to the returned string.
+
+    This is done to prevent issues caused in cases like: truncate_string(fg.blue + hello + fg.reset, 3)
+    In this case, "hello" would be truncated before fg.reset resets the color from blue. Appending the remaining style
+    sequences makes sure the style is in the same state had the entire string been printed.
 
     :param line: text to truncate
     :param max_width: the maximum display width the resulting string is allowed to have
@@ -855,6 +862,7 @@ def truncate_line(line: str, max_width: int, *, tab_width: int = 4) -> str:
     :raises: ValueError if text contains an unprintable character like a new line
              ValueError if max_width is less than 1
     """
+    import io
     from . import ansi
 
     # Handle tabs
@@ -866,12 +874,68 @@ def truncate_line(line: str, max_width: int, *, tab_width: int = 4) -> str:
     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]
+    if ansi.style_aware_wcswidth(line) <= max_width:
+        return line
+
+    # Find all style sequences in the line
+    styles = get_styles_in_text(line)
+
+    # Add characters one by one and preserve all style sequences
+    done = False
+    index = 0
+    total_width = 0
+    truncated_buf = io.StringIO()
+
+    while not done:
+        # Check if a style sequence is at this index. These don't count toward display width.
+        if index in styles:
+            truncated_buf.write(styles[index])
+            style_len = len(styles[index])
+            styles.pop(index)
+            index += style_len
+            continue
+
+        char = line[index]
+        char_width = ansi.style_aware_wcswidth(char)
+
+        # This char will make the text too wide, add the ellipsis instead
+        if char_width + total_width >= max_width:
+            char = constants.HORIZONTAL_ELLIPSIS
+            char_width = ansi.style_aware_wcswidth(char)
+            done = True
+
+        total_width += char_width
+        truncated_buf.write(char)
+        index += 1
+
+    # Append remaining style sequences from original string
+    truncated_buf.write(''.join(styles.values()))
+
+    return truncated_buf.getvalue()
+
+
+def get_styles_in_text(text: str) -> Dict[int, str]:
+    """
+    Return an OrderedDict containing all ANSI style sequences found in a string
+
+    The structure of the dictionary is:
+        key: index where sequences begins
+        value: ANSI style sequence found at index in text
+
+    Keys are in ascending order
+
+    :param text: text to search for style sequences
+    """
+    from . import ansi
+
+    start = 0
+    styles = collections.OrderedDict()
 
-        line += "\N{HORIZONTAL ELLIPSIS}"
+    while True:
+        match = ansi.ANSI_STYLE_RE.search(text, start)
+        if match is None:
+            break
+        styles[match.start()] = match.group()
+        start += len(match.group())
 
-    return line
+    return styles

tests/test_utils.py

@@ -10,6 +10,7 @@
 import pytest
 
 import cmd2.utils as cu
+from cmd2.constants import HORIZONTAL_ELLIPSIS
 
 HELLO_WORLD = 'Hello, world!'
 
@@ -297,7 +298,13 @@ def test_truncate_line():
     line = 'long'
     max_width = 3
     truncated = cu.truncate_line(line, max_width)
-    assert truncated == 'lo\N{HORIZONTAL ELLIPSIS}'
+    assert truncated == 'lo' + HORIZONTAL_ELLIPSIS
+
+def test_truncate_line_already_fits():
+    line = 'long'
+    max_width = 4
+    truncated = cu.truncate_line(line, max_width)
+    assert truncated == line
 
 def test_truncate_line_with_newline():
     line = 'fo\no'
@@ -315,20 +322,44 @@ def test_truncate_line_wide_text():
     line = '苹苹other'
     max_width = 6
     truncated = cu.truncate_line(line, max_width)
-    assert truncated == '苹苹o\N{HORIZONTAL ELLIPSIS}'
+    assert truncated == '苹苹o' + HORIZONTAL_ELLIPSIS
 
 def test_truncate_line_split_wide_text():
     """Test when truncation results in a string which is shorter than max_width"""
     line = '1苹2苹'
     max_width = 3
     truncated = cu.truncate_line(line, max_width)
-    assert truncated == '1\N{HORIZONTAL ELLIPSIS}'
+    assert truncated == '1' + HORIZONTAL_ELLIPSIS
 
 def test_truncate_line_tabs():
     line = 'has\ttab'
     max_width = 9
     truncated = cu.truncate_line(line, max_width)
-    assert truncated == 'has    t\N{HORIZONTAL ELLIPSIS}'
+    assert truncated == 'has    t' + HORIZONTAL_ELLIPSIS
+
+def test_truncate_with_style():
+    from cmd2 import ansi
+
+    before_style = ansi.fg.blue + ansi.UNDERLINE_ENABLE
+    after_style = ansi.fg.reset + ansi.UNDERLINE_DISABLE
+
+    # Style only before truncated text
+    line = before_style + 'long'
+    max_width = 3
+    truncated = cu.truncate_line(line, max_width)
+    assert truncated == before_style + 'lo' + HORIZONTAL_ELLIPSIS
+
+    # Style before and after truncated text
+    line = before_style + 'long' + after_style
+    max_width = 3
+    truncated = cu.truncate_line(line, max_width)
+    assert truncated == before_style + 'lo' + HORIZONTAL_ELLIPSIS + after_style
+
+    # Style only after truncated text
+    line = 'long' + after_style
+    max_width = 3
+    truncated = cu.truncate_line(line, max_width)
+    assert truncated == 'lo' + HORIZONTAL_ELLIPSIS + after_style
 
 def test_align_text_fill_char_is_tab():
     text = 'foo'
@@ -337,6 +368,15 @@ def test_align_text_fill_char_is_tab():
     aligned = cu.align_text(text, cu.TextAlignment.LEFT, fill_char=fill_char, width=width)
     assert aligned == text + '  '
 
+def test_align_text_fill_char_has_color():
+    from cmd2 import ansi
+
+    text = 'foo'
+    fill_char = ansi.fg.bright_yellow + '-' + ansi.fg.reset
+    width = 5
+    aligned = cu.align_text(text, cu.TextAlignment.LEFT, fill_char=fill_char, width=width)
+    assert aligned == text + fill_char * 2
+
 def test_align_text_width_is_too_small():
     text = 'foo'
     fill_char = '-'
@@ -351,7 +391,7 @@ def test_align_text_fill_char_is_too_long():
     with pytest.raises(TypeError):
         cu.align_text(text, cu.TextAlignment.LEFT, fill_char=fill_char, width=width)
 
-def test_align_text_fill_char_is_unprintable():
+def test_align_text_fill_char_is_newline():
     text = 'foo'
     fill_char = '\n'
     width = 5
@@ -384,15 +424,15 @@ def test_align_text_wider_than_width_truncate():
     fill_char = '-'
     width = 8
     aligned = cu.align_text(text, cu.TextAlignment.LEFT, fill_char=fill_char, width=width, truncate=True)
-    assert aligned == 'long te\N{HORIZONTAL ELLIPSIS}'
+    assert aligned == 'long te' + HORIZONTAL_ELLIPSIS
 
 def test_align_text_wider_than_width_truncate_add_fill():
     """Test when truncation results in a string which is shorter than width and align_text adds filler"""
     text = '1苹2苹'
     fill_char = '-'
     width = 3
     aligned = cu.align_text(text, cu.TextAlignment.LEFT, fill_char=fill_char, width=width, truncate=True)
-    assert aligned == '1\N{HORIZONTAL ELLIPSIS}-'
+    assert aligned == '1' + HORIZONTAL_ELLIPSIS + fill_char
 
 def test_align_text_has_unprintable():
     text = 'foo\x02'