tbl: add _ColumnCellCollection magic methods

Steve Canny · Steve Canny · commit 57516ae43d01 · 2013-12-22T16:27:42.000-08:00
diff --git a/docx/table.py b/docx/table.py
@@ -91,6 +91,35 @@ def __init__(self, tbl, gridCol):
         self._tbl = tbl
         self._gridCol = gridCol
 
+    def __getitem__(self, idx):
+        """
+        Provide indexed access, (e.g. 'cells[0]')
+        """
+        try:
+            tr = self._tr_lst[idx]
+        except IndexError:
+            msg = "cell index [%d] is out of range" % idx
+            raise IndexError(msg)
+        tc = tr.tc_lst[self._col_idx]
+        return _Cell(tc)
+
+    def __iter__(self):
+        for tr in self._tr_lst:
+            tc = tr.tc_lst[self._col_idx]
+            yield _Cell(tc)
+
+    def __len__(self):
+        return len(self._tr_lst)
+
+    @property
+    def _col_idx(self):
+        gridCol_lst = self._tbl.tblGrid.gridCol_lst
+        return gridCol_lst.index(self._gridCol)
+
+    @property
+    def _tr_lst(self):
+        return self._tbl.tr_lst
+
 
 class _ColumnCollection(object):
     """
diff --git a/features/tbl-item-access.feature b/features/tbl-item-access.feature
@@ -23,7 +23,6 @@ Feature: Access table rows, columns, and cells
       Then I can iterate over the column collection
        And I can access a collection column by index
 
-  @wip
   Scenario: Access cell collection of table column
      Given a table column having two cells
       Then I can access the cell collection of the column
@@ -34,7 +33,6 @@ Feature: Access table rows, columns, and cells
       Then I can access the cell collection of the row
        And I can get the length of the row cell collection
 
-  @wip
   Scenario: Access cell in column cell collection
      Given a column cell collection having two cells
       Then I can iterate over the column cells
diff --git a/tests/test_table.py b/tests/test_table.py
@@ -84,6 +84,46 @@ def column(self):
         return _Column(None, None)
 
 
+class Describe_ColumnCellCollection(object):
+
+    def it_knows_how_many_cells_it_contains(self, cells_fixture):
+        cells, cell_count = cells_fixture
+        assert len(cells) == cell_count
+
+    def it_can_iterate_over_its__Cell_instances(self, cells_fixture):
+        cells, cell_count = cells_fixture
+        actual_count = 0
+        for cell in cells:
+            assert isinstance(cell, _Cell)
+            actual_count += 1
+        assert actual_count == cell_count
+
+    def it_provides_indexed_access_to_cells(self, cells_fixture):
+        cells, cell_count = cells_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, cells_fixture):
+        cells, cell_count = cells_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 cells_fixture(self):
+        cell_count = 2
+        tbl = _tbl_bldr(rows=cell_count, cols=1).element
+        gridCol = tbl.tblGrid.gridCol_lst[0]
+        cells = _ColumnCellCollection(tbl, gridCol)
+        return cells, cell_count
+
+
 class Describe_ColumnCollection(object):
 
     def it_knows_how_many_columns_it_contains(self, columns_fixture):
@@ -141,6 +181,10 @@ def cells_access_fixture(self):
 
 class Describe_RowCellCollection(object):
 
+    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_can_iterate_over_its__Cell_instances(self, cell_count_fixture):
         cells, cell_count = cell_count_fixture
         actual_count = 0
@@ -149,10 +193,6 @@ def it_can_iterate_over_its__Cell_instances(self, cell_count_fixture):
             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):
