tbl: add _RowCollection magic methods

Steve Canny · Steve Canny · commit f79e7f224a18 · 2013-12-21T23:27:49.000-08:00
diff --git a/docx/table.py b/docx/table.py
@@ -67,3 +67,20 @@ class _RowCollection(object):
     def __init__(self, tbl):
         super(_RowCollection, self).__init__()
         self._tbl = tbl
+
+    def __getitem__(self, idx):
+        """
+        Provide indexed access, (e.g. 'rows[0]')
+        """
+        try:
+            tr = self._tbl.tr_lst[idx]
+        except IndexError:
+            msg = "row index [%d] out of range" % idx
+            raise IndexError(msg)
+        return _Row(tr)
+
+    def __iter__(self):
+        return iter([_Row(tr) for tr in self._tbl.tr_lst])
+
+    def __len__(self):
+        return len(self._tbl.tr_lst)
diff --git a/features/steps/table.py b/features/steps/table.py
@@ -25,6 +25,15 @@ def given_a_table_having_two_rows(context):
 
 # then =====================================================
 
+@then('I can access the rows by index')
+def then_can_access_rows_by_index(context):
+    rows = context.table_.rows
+    row_count = len(rows)
+    for idx in range(row_count):
+        row = rows[idx]
+        assert isinstance(row, _Row)
+
+
 @then('the length of its row collection is 2')
 def then_len_of_row_collection_is_2(context):
     rows = context.table_.rows
diff --git a/features/tbl-item-access.feature b/features/tbl-item-access.feature
@@ -7,3 +7,7 @@ Feature: Access table rows, columns, and cells
      Given a table having two rows
       Then the length of its row collection is 2
        And each item in its row collection is a table row
+
+  Scenario: Access table row by collection index
+     Given a table having two rows
+      Then I can access the rows by index
diff --git a/tests/test_table.py b/tests/test_table.py
@@ -56,6 +56,46 @@ def row_access_fixture(self):
         return table
 
 
+class Describe_RowCollection(object):
+
+    def it_contains__Row_instances(self, row_count_fixture):
+        table, row_count = row_count_fixture
+        actual_count = 0
+        for row in table.rows:
+            assert isinstance(row, _Row)
+            actual_count += 1
+        assert actual_count == row_count
+
+    def it_knows_how_many_rows_it_contains(self, row_count_fixture):
+        table, row_count = row_count_fixture
+        rows = table.rows
+        assert len(rows) == row_count
+
+    def it_provides_indexed_access_to_rows(self, row_count_fixture):
+        table, row_count = row_count_fixture
+        for idx in range(-row_count, row_count):
+            row = table.rows[idx]
+            assert isinstance(row, _Row)
+
+    def it_raises_on_indexed_access_out_of_range(self, row_count_fixture):
+        table, row_count = row_count_fixture
+        with pytest.raises(IndexError):
+            too_low = -1 - row_count
+            table.rows[too_low]
+        with pytest.raises(IndexError):
+            too_high = row_count
+            table.rows[too_high]
+
+    # fixtures -------------------------------------------------------
+
+    @pytest.fixture
+    def row_count_fixture(self):
+        row_count = 2
+        tbl = _tbl_bldr(rows=row_count, cols=2).element
+        table = Table(tbl)
+        return table, row_count
+
+
 # fixtures -----------------------------------------------------------
 
 def _tbl_bldr(rows, cols):
