WIP: Try to fix tabulate._CustomTextWrap._handle_long_word breaking up ANSI escape codes.

Author: devdanzin

tabulate/__init__.py

@@ -2540,10 +2540,18 @@ def _handle_long_word(self, reversed_chunks, cur_line, cur_len, width):
             # take each charcter's width into account
             chunk = reversed_chunks[-1]
             i = 1
-            while self._len(chunk[:i]) <= space_left:
+            while len(_strip_ansi(chunk)[:i]) <= space_left:
                 i = i + 1
-            cur_line.append(chunk[: i - 1])
-            reversed_chunks[-1] = chunk[i - 1 :]
+            # Consider escape codes when breaking words up
+            total_escape_len = 0
+            if _ansi_codes.search(chunk) is not None:
+                for group, _, _, _ in _ansi_codes.findall(chunk):
+                    escape_len = len(group)
+                    # FIXME: Needs to keep track of found groups and search from there
+                    if group in chunk[: i  + total_escape_len + escape_len - 1]:
+                        total_escape_len += escape_len
+            cur_line.append(chunk[: i  + total_escape_len - 1])
+            reversed_chunks[-1] = chunk[i + total_escape_len - 1 :]
 
         # Otherwise, we have to preserve the long word intact.  Only add
         # it to the current line if there's nothing already there --

test/test_textwrapper.py

@@ -3,7 +3,7 @@
 
 import datetime
 
-from tabulate import _CustomTextWrap as CTW, tabulate
+from tabulate import _CustomTextWrap as CTW, tabulate, _strip_ansi
 from textwrap import TextWrapper as OTW
 
 from common import skip, assert_equal
@@ -158,6 +158,42 @@ def test_wrap_color_line_splillover():
     assert_equal(expected, result)
 
 
+def test_wrap_color_line_longword():
+    """TextWrapper: Wrap a line - preserve internal color tags and wrap them to
+    other lines when required, requires adding the colors tags to other lines as appropriate
+    and avoiding splitting escape codes."""
+    data = "This_is_a_\033[31mtest_string_for_testing_TextWrap\033[0m_with_colors"
+
+    expected = [
+        "This_is_a_\033[31mte\033[0m",
+        "\033[31mst_string_fo\033[0m",
+        "\033[31mr_testing_Te\033[0m",
+        "\033[31mxtWrap\033[0m_with_",
+        "colors",
+    ]
+    wrapper = CTW(width=12)
+    result = wrapper.wrap(data)
+    assert_equal(expected, result)
+
+
+def test_wrap_color_line_multiple_escapes():
+    data = '012345(\x1b[32ma\x1b[0mbc\x1b[32mdefghij\x1b[0m)'
+    expected = [
+        "012345(\x1b[32ma\x1b[0mbc\x1b[32mdefg\x1b[0m",
+        "\x1b[32mhij\x1b[0m)",
+    ]
+    wrapper = CTW(width=10)
+    result = wrapper.wrap(data)
+    assert_equal(expected, result)
+    clean_data = _strip_ansi(data)
+    for width in range(2, len(clean_data)):
+        # Currently fails with 14, 15 and 16, because a escape code gets split at the end
+        wrapper = CTW(width=width)
+        result = wrapper.wrap(data)
+        # print(width, result)
+        # Comparing after stripping ANSI should be enough to catch broken escape codes
+        assert_equal(clean_data, _strip_ansi("".join(result)))
+
 def test_wrap_datetime():
     """TextWrapper: Show that datetimes can be wrapped without crashing"""
     data = [