bmk: add Bookmarks.__getitem__() and .__iter__()

Benjamin Toornstra · scanny · commit bd58f434117f · 2018-12-27T19:56:10.000-08:00
Adding .__iter__() is not strictly required to enable iteration, but it
improves the performance of iteration significantly by avoiding the
default implementation which would repeatedly parse the bookmark pairs
in the document.
diff --git a/docx/bookmark.py b/docx/bookmark.py
@@ -6,18 +6,34 @@
     absolute_import, division, print_function, unicode_literals
 )
 
+from collections import Sequence
 from itertools import chain
 
 from docx.oxml.ns import qn
 from docx.shared import lazyproperty
 
 
-class Bookmarks(object):
-    """Sequence of |Bookmark| objects."""
+class Bookmarks(Sequence):
+    """Sequence of |Bookmark| objects.
+
+    Supports indexed access (including slices), `len()`, and iteration. Iteration will
+    perform significantly better than repeated indexed access.
+    """
 
     def __init__(self, document_part):
         self._document_part = document_part
 
+    def __getitem__(self, idx):
+        """Supports indexed and sliced access."""
+        bookmark_pairs = self._finder.bookmark_pairs
+        if isinstance(idx, slice):
+            return [_Bookmark(pair) for pair in bookmark_pairs[idx]]
+        return _Bookmark(bookmark_pairs[idx])
+
+    def __iter__(self):
+        """Supports iteration."""
+        return (_Bookmark(pair) for pair in self._finder.bookmark_pairs)
+
     def __len__(self):
         return len(self._finder.bookmark_pairs)
 
@@ -27,6 +43,13 @@ def _finder(self):
         return _DocumentBookmarkFinder(self._document_part)
 
 
+class _Bookmark(object):
+    """Proxy for a (w:bookmarkStart, w:bookmarkEnd) element pair."""
+
+    def __init__(self, bookmark_pair):
+        self._bookmarkStart, self._bookmarkEnd = bookmark_pair
+
+
 class _DocumentBookmarkFinder(object):
     """Provides access to bookmark oxml elements in an overall document."""
 
diff --git a/features/bmk-bookmarks.feature b/features/bmk-bookmarks.feature
@@ -4,7 +4,6 @@ Feature: Access a bookmark
   I need sequence operations on Bookmarks
 
 
-  @wip
   Scenario: Bookmarks is a sequence
     Given a Bookmarks object of length 5 as bookmarks
      Then len(bookmarks) == 5
diff --git a/tests/test_bookmark.py b/tests/test_bookmark.py
@@ -6,7 +6,12 @@
 
 import pytest
 
-from docx.bookmark import Bookmarks, _DocumentBookmarkFinder, _PartBookmarkFinder
+from docx.bookmark import (
+    _Bookmark,
+    Bookmarks,
+    _DocumentBookmarkFinder,
+    _PartBookmarkFinder,
+)
 from docx.opc.part import Part, XmlPart
 from docx.parts.document import DocumentPart
 
@@ -24,9 +29,59 @@
 
 class DescribeBookmarks(object):
 
-    def it_knows_how_many_bookmarks_the_document_contains(
-        self, _finder_prop_, finder_
+    def it_provides_access_to_bookmarks_by_index(
+        self, _finder_prop_, finder_, _Bookmark_, bookmark_
+    ):
+        bookmarkStarts = tuple(element("w:bookmarkStart") for _ in range(3))
+        bookmarkEnds = tuple(element("w:bookmarkEnd") for _ in range(3))
+        _finder_prop_.return_value = finder_
+        finder_.bookmark_pairs = zip(bookmarkStarts, bookmarkEnds)
+        _Bookmark_.return_value = bookmark_
+        bookmarks = Bookmarks(None)
+
+        bookmark = bookmarks[1]
+
+        _Bookmark_.assert_called_once_with((bookmarkStarts[1], bookmarkEnds[1]))
+        assert bookmark == bookmark_
+
+    def it_provides_access_to_bookmarks_by_slice(
+        self, _finder_prop_, finder_, _Bookmark_, bookmark_
+    ):
+        bookmarkStarts = tuple(element("w:bookmarkStart") for _ in range(4))
+        bookmarkEnds = tuple(element("w:bookmarkEnd") for _ in range(4))
+        _finder_prop_.return_value = finder_
+        finder_.bookmark_pairs = zip(bookmarkStarts, bookmarkEnds)
+        _Bookmark_.return_value = bookmark_
+        bookmarks = Bookmarks(None)
+
+        bookmarks_slice = bookmarks[1:3]
+
+        assert _Bookmark_.call_args_list == [
+            call((bookmarkStarts[1], bookmarkEnds[1])),
+            call((bookmarkStarts[2], bookmarkEnds[2])),
+        ]
+        assert bookmarks_slice == [bookmark_, bookmark_]
+
+    def it_can_iterate_its_bookmarks(
+        self, _finder_prop_, finder_, _Bookmark_, bookmark_
     ):
+        bookmarkStarts = tuple(element("w:bookmarkStart") for _ in range(3))
+        bookmarkEnds = tuple(element("w:bookmarkEnd") for _ in range(3))
+        _finder_prop_.return_value = finder_
+        finder_.bookmark_pairs = zip(bookmarkStarts, bookmarkEnds)
+        _Bookmark_.return_value = bookmark_
+        bookmarks = Bookmarks(None)
+
+        _bookmarks = list(b for b in bookmarks)
+
+        assert _Bookmark_.call_args_list == [
+            call((bookmarkStarts[0], bookmarkEnds[0])),
+            call((bookmarkStarts[1], bookmarkEnds[1])),
+            call((bookmarkStarts[2], bookmarkEnds[2])),
+        ]
+        assert _bookmarks == [bookmark_, bookmark_, bookmark_]
+
+    def it_knows_how_many_bookmarks_the_document_contains(self, _finder_prop_, finder_):
         _finder_prop_.return_value = finder_
         finder_.bookmark_pairs = tuple((1, 2) for _ in range(42))
         bookmarks = Bookmarks(None)
@@ -48,6 +103,14 @@ def it_provides_access_to_its_bookmark_finder_to_help(
 
     # fixture components ---------------------------------------------
 
+    @pytest.fixture
+    def _Bookmark_(self, request):
+        return class_mock(request, 'docx.bookmark._Bookmark')
+
+    @pytest.fixture
+    def bookmark_(self, request):
+        return instance_mock(request, _Bookmark)
+
     @pytest.fixture
     def _DocumentBookmarkFinder_(self, request):
         return class_mock(request, 'docx.bookmark._DocumentBookmarkFinder')
