bmk: add BlockItemContainer.end_bookmark()

scanny · scanny · commit 99778bb2045c · 2021-05-15T15:32:43.000-07:00
diff --git a/docx/blkcntnr.py b/docx/blkcntnr.py
@@ -54,7 +54,9 @@ def add_table(self, rows, cols, width):
 
     def end_bookmark(self, bookmark):
         """Return `bookmark` after closing it after last block item in container."""
-        raise NotImplementedError
+        if bookmark.is_closed:
+            raise ValueError("bookmark already closed")
+        return bookmark.close(self._element.add_bookmarkEnd(bookmark.id))
 
     @property
     def paragraphs(self):
diff --git a/docx/bookmark.py b/docx/bookmark.py
@@ -76,11 +76,24 @@ class _Bookmark(object):
     def __init__(self, bookmark_pair):
         self._bookmarkStart, self._bookmarkEnd = bookmark_pair
 
+    def close(self, bookmarkEnd):
+        """Return self after setting end marker to `bookmarkEnd`.
+
+        Raises ValueError if this bookmark is already closed or if `id` attribute of
+        `bookmarkEnd` does not match that of the `bookmarkStart` element.
+        """
+        raise NotImplementedError
+
     @property
     def id(self):
         """Provides access to the bookmark id."""
         return self._bookmarkStart.id
 
+    @property
+    def is_closed(self):
+        """True if this bookmark has both a start and end element."""
+        raise NotImplementedError
+
     @property
     def name(self):
         """Provides access to the bookmark name."""
diff --git a/docx/oxml/document.py b/docx/oxml/document.py
@@ -1,9 +1,6 @@
 # encoding: utf-8
 
-"""
-Custom element classes that correspond to the document part, e.g.
-<w:document>.
-"""
+"""Custom element classes that correspond to the document part, e.g. <w:document>."""
 
 from .xmlchemy import BaseOxmlElement, ZeroOrOne, ZeroOrMore
 
@@ -30,8 +27,20 @@ class CT_Body(BaseOxmlElement):
     p = ZeroOrMore("w:p", successors=("w:sectPr",))
     tbl = ZeroOrMore("w:tbl", successors=("w:sectPr",))
     bookmarkStart = ZeroOrMore("w:bookmarkStart", successors=("w:sectPr",))
+    bookmarkEnd = ZeroOrMore("w:bookmarkEnd", successors=("w:sectPr",))
     sectPr = ZeroOrOne("w:sectPr", successors=())
 
+    def add_bookmarkEnd(self, bookmark_id):
+        """Return `w:bookmarkEnd` element added at end of document.
+
+        The newly added `w:bookmarkEnd` element is linked to it's `w:bookmarkStart`
+        counterpart by `bookmark_id`. It is the caller's responsibility to determine
+        `bookmark_id` matches that of the intended `bookmarkStart` element.
+        """
+        bookmarkEnd = self._add_bookmarkEnd()
+        bookmarkEnd.id = bookmark_id
+        return bookmarkEnd
+
     def add_bookmarkStart(self, name, bookmark_id):
         """Return `w:bookmarkStart` element added at end of document.
 
diff --git a/docx/oxml/section.py b/docx/oxml/section.py
@@ -23,6 +23,18 @@ class CT_HdrFtr(BaseOxmlElement):
     p = ZeroOrMore("w:p", successors=())
     tbl = ZeroOrMore("w:tbl", successors=())
     bookmarkStart = ZeroOrMore("w:bookmarkStart", successors=())
+    bookmarkEnd = ZeroOrMore("w:bookmarkEnd", successors=())
+
+    def add_bookmarkEnd(self, bookmark_id):
+        """Return `w:bookmarkEnd` element added at end of this header or footer.
+
+        The newly added `w:bookmarkEnd` element is linked to it's `w:bookmarkStart`
+        counterpart by `bookmark_id`. It is the caller's responsibility to determine
+        `bookmark_id` matches that of the intended `bookmarkStart` element.
+        """
+        bookmarkEnd = self._add_bookmarkEnd()
+        bookmarkEnd.id = bookmark_id
+        return bookmarkEnd
 
     def add_bookmarkStart(self, name, bookmark_id):
         """Return `w:bookmarkStart` element added at the end of this header or footer.
diff --git a/docx/oxml/table.py b/docx/oxml/table.py
@@ -399,6 +399,18 @@ class CT_Tc(BaseOxmlElement):
     p = OneOrMore("w:p")
     tbl = OneOrMore("w:tbl")
     bookmarkStart = ZeroOrMore("w:bookmarkStart", successors=())
+    bookmarkEnd = ZeroOrMore("w:bookmarkEnd", successors=())
+
+    def add_bookmarkEnd(self, bookmark_id):
+        """Return `w:bookmarkEnd` element added at end of this cell.
+
+        The newly added `w:bookmarkEnd` element is linked to it's `w:bookmarkStart`
+        counterpart by `bookmark_id`. It is the caller's responsibility to determine
+        `bookmark_id` matches that of the intended `bookmarkStart` element.
+        """
+        bookmarkEnd = self._add_bookmarkEnd()
+        bookmarkEnd.id = bookmark_id
+        return bookmarkEnd
 
     def add_bookmarkStart(self, name, bookmark_id):
         """Return `w:bookmarkStart` element added at the end of this cell.
diff --git a/tests/test_blkcntnr.py b/tests/test_blkcntnr.py
@@ -50,6 +50,29 @@ def it_can_add_a_table(self, add_table_fixture):
         assert table._element.xml == expected_xml
         assert table._parent is blkcntnr
 
+    def it_can_end_a_bookmark(self, end_bookmark_fixture, bookmark_):
+        blockContainer, bookmark_id, expected_xml = end_bookmark_fixture
+        bookmark_.close.return_value = bookmark_
+        bookmark_.id = bookmark_id
+        bookmark_.is_closed = False
+        blkcntnr = BlockItemContainer(blockContainer, None)
+
+        bookmark = blkcntnr.end_bookmark(bookmark_)
+
+        bookmark_.close.assert_called_once_with(
+            blockContainer.xpath("w:bookmarkEnd")[-1]
+        )
+        assert blkcntnr._element.xml == expected_xml
+        assert bookmark is bookmark_
+
+    def but_it_raises_when_bookmark_is_already_closed(self, bookmark_):
+        bookmark_.is_closed = True
+        blkcntnr = BlockItemContainer(None, None)
+
+        with pytest.raises(ValueError) as e:
+            blkcntnr.end_bookmark(bookmark_)
+        assert "bookmark already closed" in str(e.value)
+
     def it_provides_access_to_the_paragraphs_it_contains(self, paragraphs_fixture):
         # ---test len(), iterable, and indexed access---
         blkcntnr, expected_count = paragraphs_fixture
@@ -160,6 +183,24 @@ def bookmarks_fixture(self, request):
         parent_part_ = instance_mock(request, PartCls)
         return parent_part_
 
+    @pytest.fixture(
+        params=[
+            # ---document body---
+            ("w:body", 0, "w:body/w:bookmarkEnd{w:id=0}"),
+            # ---table cell---
+            ("w:tc/w:p", 1, "w:tc/(w:p,w:bookmarkEnd{w:id=1})"),
+            # ---header---
+            ("w:hdr/w:p", 42, "w:hdr/(w:p,w:bookmarkEnd{w:id=42})"),
+            # ---footer---
+            ("w:ftr/w:p", 24, "w:ftr/(w:p,w:bookmarkEnd{w:id=24})"),
+        ]
+    )
+    def end_bookmark_fixture(self, request):
+        cxml, bookmark_id, expected_cxml = request.param
+        blockContainer = element(cxml)
+        expected_xml = xml(expected_cxml)
+        return blockContainer, bookmark_id, expected_xml
+
     @pytest.fixture(
         params=[
             ("w:body", 0),
