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

Benjamin Toornstra · scanny · commit 425245c84322 · 2019-07-13T16:34:35.000-07: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
@@ -4,18 +4,34 @@
 
 from __future__ import 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)
 
@@ -25,6 +41,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
 
@@ -23,6 +28,60 @@
 
 
 class DescribeBookmarks(object):
+    """Unit-test suite for `docx.bookmark.Bookmarks` object."""
+
+    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))
@@ -45,6 +104,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")
@@ -125,9 +192,7 @@ def it_provides_an_iter_start_end_pairs_interface_method(
         assert pairs == _iter_start_end_pairs_.return_value
 
     def it_gathers_all_the_bookmark_start_and_end_elements_to_help(self, part_):
-        body = element(
-            "w:body/(w:bookmarkStart,w:p,w:bookmarkEnd,w:p,w:bookmarkStart)"
-        )
+        body = element("w:body/(w:bookmarkStart,w:p,w:bookmarkEnd,w:p,w:bookmarkStart)")
         part_.element = body
         finder = _PartBookmarkFinder(part_)
 
@@ -190,7 +255,9 @@ def it_iterates_bookmarkStart_elements_to_help(self, _all_starts_and_ends_prop_)
         starts = list(finder._iter_starts())
 
         assert starts == [
-            (0, starts_and_ends[0]), (2, starts_and_ends[2]), (4, starts_and_ends[4])
+            (0, starts_and_ends[0]),
+            (2, starts_and_ends[2]),
+            (4, starts_and_ends[4]),
         ]
 
     def it_finds_the_matching_end_for_a_start_to_help(
@@ -279,7 +346,7 @@ def name_used_fixture(self, request):
 
     @pytest.fixture
     def _all_starts_and_ends_prop_(self, request):
-        return property_mock(request, _PartBookmarkFinder, '_all_starts_and_ends')
+        return property_mock(request, _PartBookmarkFinder, "_all_starts_and_ends")
 
     @pytest.fixture
     def _init_(self, request):
@@ -303,7 +370,7 @@ def _name_already_used_(self, request):
 
     @pytest.fixture
     def _names_so_far_prop_(self, request):
-        return property_mock(request, _PartBookmarkFinder, '_names_so_far')
+        return property_mock(request, _PartBookmarkFinder, "_names_so_far")
 
     @pytest.fixture
     def names_so_far_(self, request):
