Align expanded tabs with terminal columns after wide characters

src/prompt_toolkit/layout/processors.py

@@ -26,6 +26,7 @@
 from prompt_toolkit.search import SearchDirection
 from prompt_toolkit.utils import to_int, to_str
 
+from .screen import _CHAR_CACHE
 from .utils import explode_text_fragments
 
 if TYPE_CHECKING:
@@ -716,28 +717,38 @@ def apply_transformation(self, ti: TransformationInput) -> Transformation:
         position_mappings = {}
         result_fragments: StyleAndTextTuples = []
         pos = 0
+        column = 0
+        source_pos = 0
 
-        for i, fragment_and_text in enumerate(fragments):
-            position_mappings[i] = pos
+        for fragment_and_text in fragments:
+            if "[ZeroWidthEscape]" in fragment_and_text[0]:
+                result_fragments.append(fragment_and_text)
+                continue
+
+            position_mappings[source_pos] = pos
 
             if fragment_and_text[1] == "\t":
                 # Calculate how many characters we have to insert.
-                count = tabstop - (pos % tabstop)
+                count = tabstop - (column % tabstop)
                 if count == 0:
                     count = tabstop
 
                 # Insert tab.
                 result_fragments.append((style, separator1))
                 result_fragments.append((style, separator2 * (count - 1)))
                 pos += count
+                column += count
             else:
                 result_fragments.append(fragment_and_text)
                 pos += 1
+                column += _CHAR_CACHE[fragment_and_text[1], fragment_and_text[0]].width
+
+            source_pos += 1
 
-        position_mappings[len(fragments)] = pos
+        position_mappings[source_pos] = pos
         # Add `pos+1` to mapping, because the cursor can be right after the
         # line as well.
-        position_mappings[len(fragments) + 1] = pos + 1
+        position_mappings[source_pos + 1] = pos + 1
 
         def source_to_display(from_position: int) -> int:
             "Maps original cursor position to the new one."

tests/test_tabs_processor.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import pytest
+
+from prompt_toolkit.document import Document
+from prompt_toolkit.formatted_text import fragment_list_to_text
+from prompt_toolkit.layout.controls import BufferControl
+from prompt_toolkit.layout.processors import TabsProcessor, TransformationInput
+
+
+@pytest.mark.parametrize(
+    "text, expected, mappings",
+    [
+        ("a\tb", "a   b", [0, 1, 4, 5]),
+        ("\x01\tb", "\x01  b", [0, 1, 3, 4]),
+        ("\x80\tb", "\x80    b", [0, 1, 5, 6]),
+        ("界\tb", "界  b", [0, 1, 3, 4]),
+        ("界界\tb", "界界    b", [0, 1, 2, 6, 7]),
+        ("e\u0301\tb", "e\u0301   b", [0, 1, 2, 5, 6]),
+        ("界\t界\tb", "界  界  b", [0, 1, 3, 4, 6, 7]),
+    ],
+)
+def test_tabs_follow_terminal_columns(text, expected, mappings):
+    result = TabsProcessor(tabstop=4, char1=" ", char2=" ").apply_transformation(
+        TransformationInput(
+            buffer_control=BufferControl(),
+            document=Document(text),
+            lineno=0,
+            source_to_display=lambda i: i,
+            fragments=[("", text)],
+            width=80,
+            height=24,
+        )
+    )
+    assert fragment_list_to_text(result.fragments) == expected
+    assert [result.source_to_display(i) for i in range(len(text) + 1)] == mappings
+    for source, display in enumerate(mappings):
+        assert result.display_to_source(display) == source
+    for source, char in enumerate(text):
+        if char == "\t":
+            for display in range(mappings[source], mappings[source + 1]):
+                assert result.display_to_source(display) == source
+
+
+@pytest.mark.parametrize("escape", ["\x1b[31m", "\x1b]0;title\tvalue\x07"])
+def test_tabs_preserve_zero_width_escapes(escape):
+    result = TabsProcessor(tabstop=4, char1=" ", char2=" ").apply_transformation(
+        TransformationInput(
+            buffer_control=BufferControl(),
+            document=Document("a\tb"),
+            lineno=0,
+            source_to_display=lambda i: i,
+            fragments=[("", "a"), ("[ZeroWidthEscape]", escape), ("", "\tb")],
+            width=80,
+            height=24,
+        )
+    )
+    assert fragment_list_to_text(result.fragments) == "a   b"
+    assert (
+        "".join(
+            text for style, text in result.fragments if "[ZeroWidthEscape]" in style
+        )
+        == escape
+    )
+    assert [result.source_to_display(i) for i in range(4)] == [0, 1, 4, 5]
+    assert [result.display_to_source(i) for i in range(6)] == [0, 1, 1, 1, 2, 3]