hdr: add _Header._add_header_part()

ondrej-111 · scanny · commit 2afbbb6b87f6 · 2019-01-06T14:26:04.000-08:00
diff --git a/docx/oxml/section.py b/docx/oxml/section.py
@@ -12,6 +12,7 @@
     BaseOxmlElement,
     OptionalAttribute,
     RequiredAttribute,
+    ZeroOrMore,
     ZeroOrOne,
 )
 
@@ -56,11 +57,22 @@ class CT_SectPr(BaseOxmlElement):
         'w:noEndnote', 'w:titlePg', 'w:textDirection', 'w:bidi', 'w:rtlGutter',
         'w:docGrid', 'w:printerSettings', 'w:sectPrChange',
     )
+    headerReference = ZeroOrMore("w:headerReference", successors=_tag_seq)
     type = ZeroOrOne("w:type", successors=_tag_seq[3:])
     pgSz = ZeroOrOne("w:pgSz", successors=_tag_seq[4:])
     pgMar = ZeroOrOne("w:pgMar", successors=_tag_seq[5:])
     del _tag_seq
 
+    def add_headerReference(self, type_, rId):
+        """Return newly added CT_HdrFtrRef element of *type_* with *rId*.
+
+        The element tag is `w:headerReference`.
+        """
+        headerReference = self._add_headerReference()
+        headerReference.type_ = type_
+        headerReference.rId = rId
+        return headerReference
+
     @property
     def bottom_margin(self):
         """
diff --git a/docx/parts/document.py b/docx/parts/document.py
@@ -24,6 +24,10 @@ class DocumentPart(XmlPart):
     objects provides access to this part object for that purpose.
     """
 
+    def add_header_part(self):
+        """Return (header_part, rId) pair for newly-created header part."""
+        raise NotImplementedError
+
     @property
     def core_properties(self):
         """
diff --git a/docx/parts/hdrftr.py b/docx/parts/hdrftr.py
@@ -0,0 +1,11 @@
+# encoding: utf-8
+
+"""Header and footer part objects"""
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from docx.opc.part import XmlPart
+
+
+class HeaderPart(XmlPart):
+    """Definition of a section header."""
diff --git a/docx/section.py b/docx/section.py
@@ -253,7 +253,9 @@ def is_linked_to_previous(self, value):
 
     def _add_header_part(self):
         """Return newly-added header part."""
-        raise NotImplementedError
+        header_part, rId = self._document_part.add_header_part()
+        self._sectPr.add_headerReference(WD_HEADER_FOOTER.PRIMARY, rId)
+        return header_part
 
     def _drop_header_part(self):
         """Remove header definition associated with this section."""
diff --git a/tests/test_section.py b/tests/test_section.py
@@ -8,6 +8,7 @@
 
 from docx.enum.section import WD_ORIENT, WD_SECTION
 from docx.parts.document import DocumentPart
+from docx.parts.hdrftr import HeaderPart
 from docx.section import _Footer, _Header, Section, Sections
 from docx.shared import Inches
 
@@ -386,6 +387,19 @@ def it_can_change_whether_it_is_linked_to_previous_header(
         assert _drop_header_part_.call_args_list == [call(header)] * drop_calls
         assert _add_header_part_.call_args_list == [call(header)] * add_calls
 
+    def it_can_add_a_header_part_to_help(self, document_part_, header_part_):
+        sectPr = element("w:sectPr{r:a=b}")
+        document_part_.add_header_part.return_value = header_part_, "rId3"
+        header = _Header(sectPr, document_part_)
+
+        header_part = header._add_header_part()
+
+        document_part_.add_header_part.assert_called_once_with()
+        assert sectPr.xml == xml(
+            "w:sectPr{r:a=b}/w:headerReference{w:type=default,r:id=rId3}"
+        )
+        assert header_part is header_part_
+
     def it_can_drop_the_related_header_part_to_help(self, document_part_):
         sectPr = element("w:sectPr{r:a=b}/w:headerReference{w:type=default,r:id=rId42}")
         header = _Header(sectPr, document_part_)
@@ -449,3 +463,7 @@ def _drop_header_part_(self, request):
     @pytest.fixture
     def _has_header_part_prop_(self, request):
         return property_mock(request, _Header, "_has_header_part")
+
+    @pytest.fixture
+    def header_part_(self, request):
+        return instance_mock(request, HeaderPart)
