From 70ce4dbbf66dbfbb435ee4b7801fcc9b2a196cd5 Mon Sep 17 00:00:00 2001
From: Darren Burns <darrenb900@gmail.com>
Date: Wed, 1 Nov 2023 11:43:44 +0000
Subject: [PATCH] Chop cells to fit to width

---
 rich/_wrap.py       |  8 ++++----
 rich/cells.py       | 35 ++++++++++++++++++++++++-----------
 tests/test_cells.py | 20 ++++++++++++++------
 3 files changed, 42 insertions(+), 21 deletions(-)

diff --git a/rich/_wrap.py b/rich/_wrap.py
index 586350c0f..f49058116 100644
--- a/rich/_wrap.py
+++ b/rich/_wrap.py
@@ -4,7 +4,7 @@
 from typing import Iterable, List, Tuple
 
 from ._loop import loop_last
-from .cells import cell_len, chop_cells
+from .cells import cell_len, fit_to_width
 
 re_word = re.compile(r"\s*\S+\s*")
 
@@ -30,8 +30,8 @@ def divide_line(text: str, width: int, fold: bool = True) -> List[int]:
         if line_position + word_length > width:
             if word_length > width:
                 if fold:
-                    chopped_words = chop_cells(
-                        word, max_size=width, position=line_position
+                    chopped_words = fit_to_width(
+                        word, available_width=width, position=line_position
                     )
                     for last, line in loop_last(chopped_words):
                         if start:
@@ -95,4 +95,4 @@ def divide_line(text: str, width: int, fold: bool = True) -> list[int]:
 
     console = Console(width=10)
     console.print("12345 abcdefghijklmnopqrstuvwyxzABCDEFGHIJKLMNOPQRSTUVWXYZ 12345")
-    print(chop_cells("abcdefghijklmnopqrstuvwxyz", 10, position=2))
+    print(fit_to_width("abcdefghijklmnopqrstuvwxyz", 10, position=2))
diff --git a/rich/cells.py b/rich/cells.py
index fd0097f42..def22ae10 100644
--- a/rich/cells.py
+++ b/rich/cells.py
@@ -121,22 +121,35 @@ def set_cell_size(text: str, total: int) -> str:
 
 # TODO: This is inefficient
 # TODO: This might not work with CWJ type characters
-def chop_cells(text: str, max_size: int, position: int = 0) -> List[str]:
-    """Break text in to equal (cell) length strings"""
+def fit_to_width(text: str, available_width: int) -> List[str]:
+    """Fit text within a cell width.
+
+    Args:
+        text: The text to fit.
+        available_width: The width available.
+
+    Returns:
+        A list of strings such that each string in the list has cell width
+        less than or equal to the available width.
+    """
     _get_character_cell_size = get_character_cell_size
-    total_size = position
     lines: List[List[str]] = [[]]
-    append = lines[-1].append
 
+    start_new_line = lines.append
+    append_to_last_line = lines[-1].append
+
+    current_line_width = 0
     for index, character in enumerate(text):
         cell_width = _get_character_cell_size(character)
-        if total_size + cell_width > max_size:
-            lines.append([character])
-            append = lines[-1].append
-            total_size = cell_width
+        char_doesnt_fit = current_line_width + cell_width > available_width
+
+        if char_doesnt_fit:
+            start_new_line([character])
+            append_to_last_line = lines[-1].append
+            current_line_width = cell_width
         else:
-            total_size += cell_width
-            append(character)
+            append_to_last_line(character)
+            current_line_width += cell_width
 
     return ["".join(line) for line in lines]
 
@@ -144,7 +157,7 @@ def chop_cells(text: str, max_size: int, position: int = 0) -> List[str]:
 if __name__ == "__main__":  # pragma: no cover
 
     print(get_character_cell_size("😽"))
-    for line in chop_cells("""这是对亚洲语言支持的测试。面对模棱两可的想法，拒绝猜测的诱惑。""", 8):
+    for line in fit_to_width("""这是对亚洲语言支持的测试。面对模棱两可的想法，拒绝猜测的诱惑。""", 8):
         print(line)
     for n in range(80, 1, -1):
         print(set_cell_size("""这是对亚洲语言支持的测试。面对模棱两可的想法，拒绝猜测的诱惑。""", n) + "|")
diff --git a/tests/test_cells.py b/tests/test_cells.py
index c7cde3716..e8ae484c8 100644
--- a/tests/test_cells.py
+++ b/tests/test_cells.py
@@ -1,5 +1,5 @@
 from rich import cells
-from rich.cells import chop_cells
+from rich.cells import fit_to_width
 
 
 def test_cell_len_long_string():
@@ -43,11 +43,19 @@ def test_set_cell_size_infinite():
         )
 
 
-def test_chop_cells():
+def test_fit_to_width():
+    """Simple example of splitting cells into lines of width 3."""
     text = "abcdefghijk"
-    assert chop_cells(text, 3) == ["abc", "def", "ghi", "jk"]
+    assert fit_to_width(text, 3) == ["abc", "def", "ghi", "jk"]
 
 
-def test_chop_cells_position():
-    text = "0123456"
-    assert chop_cells(text, 3, position=1) == ["123", "456"]
+def test_fit_to_width_double_width_boundary():
+    """The available width lies within a double-width character."""
+    text = "ありがとう"
+    assert fit_to_width(text, 3) == ["あ", "り", "が", "と", "う"]
+
+
+def test_fit_to_width_mixed_width():
+    """Mixed single and double-width characters."""
+    text = "あ1り2が3と4う56"
+    assert fit_to_width(text, 3) == ["あ1", "り2", "が3", "と4", "う5", "6"]