Fix column wrapping breaking ANSI escape codes (fixes #307)

tabulate/__init__.py

@@ -2540,10 +2540,21 @@ 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:
+            # Only count printable characters, so strip_ansi first, index later.
+            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
+            last_group = 0
+            if _ansi_codes.search(chunk) is not None:
+                for group, _, _, _ in _ansi_codes.findall(chunk):
+                    escape_len = len(group)
+                    if group in chunk[last_group: i + total_escape_len + escape_len - 1]:
+                        total_escape_len += escape_len
+                        found = _ansi_codes.search(chunk[last_group:])
+                        last_group += found.end()
+            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/common.py

@@ -3,8 +3,8 @@
 import warnings
 
 def assert_equal(expected, result):
-    print("Expected:\n%s\n" % expected)
-    print("Got:\n%s\n" % result)
+    print("Expected:\n%r\n" % expected)
+    print("Got:\n%r\n" % result)
     assert expected == result
 
 

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[32m\x1b[0m",
+        "\x1b[32mdefghij\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)):
+        wrapper = CTW(width=width)
+        result = wrapper.wrap(data)
+        # 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 = [