hdr: add _BaseHeaderFooter.body

eupharis · Steve Canny · commit 89d74935b740 · 2016-06-07T19:15:23.000-07:00
diff --git a/docx/header.py b/docx/header.py
@@ -8,7 +8,7 @@
     absolute_import, division, print_function, unicode_literals
 )
 
-from .shared import ElementProxy
+from .shared import ElementProxy, lazyproperty
 
 
 class _BaseHeaderFooter(ElementProxy):
@@ -23,6 +23,16 @@ def __init__(self, element, parent, type):
         self._sectPr = element
         self._type = type
 
+    @lazyproperty
+    def body(self):
+        """
+        BlockItemContainer instance with contents of Header
+        """
+        headerReference = self._sectPr.get_headerReference_of_type(self._type)
+        if headerReference is None:
+            return None
+        return self.part.related_hdrftr_body(headerReference.rId)
+
     @property
     def is_linked_to_previous(self):
         """
@@ -39,3 +49,10 @@ class Header(_BaseHeaderFooter):
     """
     One of the page headers for a section.
     """
+
+
+class HeaderFooterBody(object):
+    """
+    The rich-text body of a header or footer. Supports the same rich text
+    operations as a document, such as paragraphs and tables.
+    """
diff --git a/docx/oxml/__init__.py b/docx/oxml/__init__.py
@@ -86,11 +86,14 @@ def OxmlElement(nsptag_str, attrs=None, nsdecls=None):
 register_element_cls('w:numbering',     CT_Numbering)
 register_element_cls('w:startOverride', CT_DecimalNumber)
 
-from .section import CT_PageMar, CT_PageSz, CT_SectPr, CT_SectType
-register_element_cls('w:pgMar',  CT_PageMar)
-register_element_cls('w:pgSz',   CT_PageSz)
-register_element_cls('w:sectPr', CT_SectPr)
-register_element_cls('w:type',   CT_SectType)
+from .section import (
+    CT_HdrFtrRef, CT_PageMar, CT_PageSz, CT_SectPr, CT_SectType
+)
+register_element_cls('w:headerReference', CT_HdrFtrRef)
+register_element_cls('w:pgMar',           CT_PageMar)
+register_element_cls('w:pgSz',            CT_PageSz)
+register_element_cls('w:sectPr',          CT_SectPr)
+register_element_cls('w:type',            CT_SectType)
 
 from .shape import (
     CT_Blip, CT_BlipFillProperties, CT_GraphicalObject,
diff --git a/docx/oxml/section.py b/docx/oxml/section.py
@@ -10,12 +10,23 @@
 
 from ..enum.header import WD_HEADER_FOOTER
 from ..enum.section import WD_ORIENTATION, WD_SECTION_START
-from .simpletypes import ST_SignedTwipsMeasure, ST_TwipsMeasure
+from .simpletypes import (
+    ST_RelationshipId, ST_SignedTwipsMeasure, ST_TwipsMeasure
+)
 from .xmlchemy import (
-    BaseOxmlElement, OptionalAttribute, ZeroOrMore, ZeroOrOne
+    BaseOxmlElement, OptionalAttribute, RequiredAttribute, ZeroOrMore,
+    ZeroOrOne
 )
 
 
+class CT_HdrFtrRef(BaseOxmlElement):
+    """
+    `w:headerReference` and `w:footerReference` elements, specifying the
+    various headers and footers for a section.
+    """
+    rId = RequiredAttribute('r:id', ST_RelationshipId)
+
+
 class CT_PageMar(BaseOxmlElement):
     """
     ``<w:pgMar>`` element, defining page margins.
diff --git a/docx/parts/document.py b/docx/parts/document.py
@@ -121,6 +121,13 @@ def numbering_part(self):
             self.relate_to(numbering_part, RT.NUMBERING)
             return numbering_part
 
+    def related_hdrftr_body(self, rId):
+        """
+        Return the |HeaderFooterBody| object corresponding to the related
+        part identified by *rId*.
+        """
+        raise NotImplementedError
+
     def save(self, path_or_stream):
         """
         Save this document to *path_or_stream*, which can be either a path to
diff --git a/tests/test_header.py b/tests/test_header.py
@@ -11,9 +11,11 @@
 import pytest
 
 from docx.enum.header import WD_HEADER_FOOTER
-from docx.header import Header
+from docx.header import _BaseHeaderFooter, Header, HeaderFooterBody
+from docx.parts.document import DocumentPart
 
 from .unitutil.cxml import element
+from .unitutil.mock import call, instance_mock, property_mock
 
 
 class Describe_BaseHeaderFooter(object):
@@ -22,8 +24,26 @@ def it_knows_whether_it_is_linked_to_previous(self, is_linked_fixture):
         header, expected_value = is_linked_fixture
         assert header.is_linked_to_previous is expected_value
 
+    def it_provides_access_to_its_body(self, body_fixture):
+        header, calls, expected_value = body_fixture
+        body = header.body
+        assert header.part.related_hdrftr_body.call_args_list == calls
+        assert body == expected_value
+
     # fixtures -------------------------------------------------------
 
+    @pytest.fixture(params=[
+        ('w:sectPr',                                             None),
+        ('w:sectPr/w:headerReference{w:type=even,r:id=rId6}',    None),
+        ('w:sectPr/w:headerReference{w:type=default,r:id=rId8}', 'rId8'),
+    ])
+    def body_fixture(self, request, body_, part_prop_, document_part_):
+        sectPr_cxml, rId = request.param
+        header = Header(element(sectPr_cxml), None, WD_HEADER_FOOTER.PRIMARY)
+        calls, expected_value = ([call(rId)], body_) if rId else ([], None)
+        document_part_.related_hdrftr_body.return_value = body_
+        return header, calls, expected_value
+
     @pytest.fixture(params=[
         ('w:sectPr',                                   True),
         ('w:sectPr/w:headerReference{w:type=default}', False),
@@ -33,3 +53,19 @@ def is_linked_fixture(self, request):
         sectPr_cxml, expected_value = request.param
         header = Header(element(sectPr_cxml), None, WD_HEADER_FOOTER.PRIMARY)
         return header, expected_value
+
+    # fixture components ---------------------------------------------
+
+    @pytest.fixture
+    def body_(self, request):
+        return instance_mock(request, HeaderFooterBody)
+
+    @pytest.fixture
+    def document_part_(self, request):
+        return instance_mock(request, DocumentPart)
+
+    @pytest.fixture
+    def part_prop_(self, request, document_part_):
+        return property_mock(
+            request, _BaseHeaderFooter, 'part', return_value=document_part_
+        )
