tbl: add _CellCollection magic methods

Steve Canny · Steve Canny · commit baceb4174c98 · 2013-12-22T01:00:35.000-08:00
diff --git a/docx/oxml/table.py b/docx/oxml/table.py
@@ -31,6 +31,13 @@ def new(cls):
         """
         return OxmlElement('w:tr')
 
+    @property
+    def tc_lst(self):
+        """
+        Sequence containing the ``<w:tc>`` child elements in this ``<w:tr>``.
+        """
+        return self.findall(qn('w:tc'))
+
     def _append_tc(self, tc):
         """
         Return *tc* after appending it to end of tc sequence.
diff --git a/docx/table.py b/docx/table.py
@@ -46,6 +46,9 @@ class _Cell(object):
     """
     Table cell
     """
+    def __init__(self, tc):
+        super(_Cell, self).__init__()
+        self._tc = tc
 
 
 class _CellCollection(object):
@@ -56,6 +59,23 @@ def __init__(self, tr):
         super(_CellCollection, self).__init__()
         self._tr = tr
 
+    def __getitem__(self, idx):
+        """
+        Provide indexed access, (e.g. 'cells[0]')
+        """
+        try:
+            tc = self._tr.tc_lst[idx]
+        except IndexError:
+            msg = "cell index [%d] is out of range" % idx
+            raise IndexError(msg)
+        return _Cell(tc)
+
+    def __iter__(self):
+        return iter([_Cell(tc) for tc in self._tr.tc_lst])
+
+    def __len__(self):
+        return len(self._tr.tc_lst)
+
 
 class _Column(object):
     """
diff --git a/features/tbl-item-access.feature b/features/tbl-item-access.feature
@@ -12,7 +12,6 @@ Feature: Access table rows, columns, and cells
      Given a table having two rows
       Then I can access the rows by index
 
-  @wip
   Scenario: Access cell collection of table row
      Given a table row having two cells
       Then I can access the cell collection of the row
diff --git a/tests/test_table.py b/tests/test_table.py
@@ -8,7 +8,9 @@
 
 import pytest
 
-from docx.table import _CellCollection, _Column, _Row, _RowCollection, Table
+from docx.table import (
+    _Cell, _CellCollection, _Column, _Row, _RowCollection, Table
+)
 
 from .oxml.unitdata.table import a_gridCol, a_tbl, a_tblGrid, a_tc, a_tr
 from .oxml.unitdata.text import a_p
@@ -56,6 +58,48 @@ def row_access_fixture(self):
         return table
 
 
+class Describe_CellCollection(object):
+
+    def it_can_iterate_over_its__Cell_instances(self, cell_count_fixture):
+        cells, cell_count = cell_count_fixture
+        actual_count = 0
+        for cell in cells:
+            assert isinstance(cell, _Cell)
+            actual_count += 1
+        assert actual_count == cell_count
+
+    def it_knows_how_many_cells_it_contains(self, cell_count_fixture):
+        cells, cell_count = cell_count_fixture
+        assert len(cells) == cell_count
+
+    def it_provides_indexed_access_to_cells(self, cell_count_fixture):
+        cells, cell_count = cell_count_fixture
+        for idx in range(-cell_count, cell_count):
+            cell = cells[idx]
+            assert isinstance(cell, _Cell)
+
+    def it_raises_on_indexed_access_out_of_range(self, cell_count_fixture):
+        cells, cell_count = cell_count_fixture
+        too_low = -1 - cell_count
+        too_high = cell_count
+        with pytest.raises(IndexError):
+            cells[too_low]
+        with pytest.raises(IndexError):
+            cells[too_high]
+
+    # fixtures -------------------------------------------------------
+
+    @pytest.fixture
+    def cell_count_fixture(self):
+        cell_count = 2
+        tr_bldr = a_tr().with_nsdecls()
+        for idx in range(cell_count):
+            tr_bldr.with_child(a_tc())
+        tr = tr_bldr.element
+        cells = _CellCollection(tr)
+        return cells, cell_count
+
+
 class Describe_Row(object):
 
     def it_provides_access_to_the_row_cells(self, cells_access_fixture):
