rfctr(table): reimplement CT_Tc.tc_at_grid_offset

scanny · scanny · commit 512f269b7560 · 2024-04-29T16:06:02.000-07:00
This method was formerly named `.tc_at_grid_col()`. New implementation
takes `CT_Tr.grid_before` into account.
diff --git a/src/docx/oxml/table.py b/src/docx/oxml/table.py
@@ -76,19 +76,26 @@ def grid_before(self) -> int:
             return 0
         return trPr.grid_before
 
-    def tc_at_grid_col(self, idx: int) -> CT_Tc:
-        """`<w:tc>` element appearing at grid column `idx`.
+    def tc_at_grid_offset(self, grid_offset: int) -> CT_Tc:
+        """The `tc` element in this tr at exact `grid offset`.
 
-        Raises |ValueError| if no `w:tc` element begins at that grid column.
+        Raises ValueError when this `w:tr` contains no `w:tc` with exact starting `grid_offset`.
         """
-        grid_col = 0
+        # -- account for omitted cells at the start of the row --
+        remaining_offset = grid_offset - self.grid_before
+
         for tc in self.tc_lst:
-            if grid_col == idx:
+            # -- We've gone past grid_offset without finding a tc, no sense searching further. --
+            if remaining_offset < 0:
+                break
+            # -- We've arrived at grid_offset, this is the `w:tc` we're looking for. --
+            if remaining_offset == 0:
                 return tc
-            grid_col += tc.grid_span
-            if grid_col > idx:
-                raise ValueError("no cell on grid column %d" % idx)
-        raise ValueError("index out of bounds")
+            # -- We're not there yet, skip forward the number of layout-grid cells this cell
+            # -- occupies.
+            remaining_offset -= tc.grid_span
+
+        raise ValueError(f"no `tc` element at grid_offset={grid_offset}")
 
     @property
     def tr_idx(self) -> int:
@@ -505,7 +512,7 @@ def merge(self, other_tc: CT_Tc) -> CT_Tc:
         element and `other_tc` as diagonal corners.
         """
         top, left, height, width = self._span_dimensions(other_tc)
-        top_tc = self._tbl.tr_lst[top].tc_at_grid_col(left)
+        top_tc = self._tbl.tr_lst[top].tc_at_grid_offset(left)
         top_tc._grow_to(width, height)
         return top_tc
 
@@ -731,15 +738,15 @@ def _tbl(self) -> CT_Tbl:
     @property
     def _tc_above(self) -> CT_Tc:
         """The `w:tc` element immediately above this one in its grid column."""
-        return self._tr_above.tc_at_grid_col(self.grid_offset)
+        return self._tr_above.tc_at_grid_offset(self.grid_offset)
 
     @property
     def _tc_below(self) -> CT_Tc | None:
         """The tc element immediately below this one in its grid column."""
         tr_below = self._tr_below
         if tr_below is None:
             return None
-        return tr_below.tc_at_grid_col(self.grid_offset)
+        return tr_below.tc_at_grid_offset(self.grid_offset)
 
     @property
     def _tr(self) -> CT_Row:
diff --git a/tests/oxml/test_table.py b/tests/oxml/test_table.py
@@ -34,19 +34,11 @@ def it_can_add_a_trPr(self, tr_cxml: str, expected_cxml: str):
         tr._add_trPr()
         assert tr.xml == xml(expected_cxml)
 
-    @pytest.mark.parametrize(
-        ("snippet_idx", "row_idx", "col_idx", "err_msg"),
-        [
-            (0, 0, 3, "index out of bounds"),
-            (1, 0, 1, "no cell on grid column 1"),
-        ],
-    )
-    def it_raises_on_tc_at_grid_col(
-        self, snippet_idx: int, row_idx: int, col_idx: int, err_msg: str
-    ):
+    @pytest.mark.parametrize(("snippet_idx", "row_idx", "col_idx"), [(0, 0, 3), (1, 0, 1)])
+    def it_raises_on_tc_at_grid_col(self, snippet_idx: int, row_idx: int, col_idx: int):
         tr = cast(CT_Tbl, parse_xml(snippet_seq("tbl-cells")[snippet_idx])).tr_lst[row_idx]
-        with pytest.raises(ValueError, match=err_msg):
-            tr.tc_at_grid_col(col_idx)
+        with pytest.raises(ValueError, match=f"no `tc` element at grid_offset={col_idx}"):
+            tr.tc_at_grid_offset(col_idx)
 
 
 class DescribeCT_Tc:
@@ -76,12 +68,12 @@ def it_can_merge_to_another_tc(
         top, left, height, width = 0, 1, 2, 3
         _span_dimensions_.return_value = top, left, height, width
         _tbl_.return_value.tr_lst = [tr_]
-        tr_.tc_at_grid_col.return_value = top_tc_
+        tr_.tc_at_grid_offset.return_value = top_tc_
 
         merged_tc = tc.merge(other_tc)
 
         _span_dimensions_.assert_called_once_with(tc, other_tc)
-        top_tr_.tc_at_grid_col.assert_called_once_with(left)
+        top_tr_.tc_at_grid_offset.assert_called_once_with(left)
         top_tc_._grow_to.assert_called_once_with(width, height)
         assert merged_tc is top_tc_
 
