hdr: add related_hdrftr_body

eupharis · eupharis · commit 8ed60eefaf8d · 2016-07-11T11:55:15.000-07:00
diff --git a/docx/__init__.py b/docx/__init__.py
@@ -12,6 +12,7 @@
 from docx.opc.parts.coreprops import CorePropertiesPart
 
 from docx.parts.document import DocumentPart
+from docx.parts.header import HeaderPart
 from docx.parts.image import ImagePart
 from docx.parts.numbering import NumberingPart
 from docx.parts.settings import SettingsPart
@@ -30,6 +31,7 @@ def part_class_selector(content_type, reltype):
 PartFactory.part_type_for[CT.WML_NUMBERING] = NumberingPart
 PartFactory.part_type_for[CT.WML_SETTINGS] = SettingsPart
 PartFactory.part_type_for[CT.WML_STYLES] = StylesPart
+PartFactory.part_type_for[CT.WML_HEADER] = HeaderPart
 
 del (
     CT, CorePropertiesPart, DocumentPart, NumberingPart, PartFactory,
diff --git a/docx/header.py b/docx/header.py
@@ -8,6 +8,7 @@
     absolute_import, division, print_function, unicode_literals
 )
 
+from .blkcntnr import BlockItemContainer
 from .shared import ElementProxy, lazyproperty
 
 
@@ -51,7 +52,7 @@ class Header(_BaseHeaderFooter):
     """
 
 
-class HeaderFooterBody(object):
+class HeaderFooterBody(BlockItemContainer):
     """
     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/parts/document.py b/docx/parts/document.py
@@ -126,7 +126,15 @@ def related_hdrftr_body(self, rId):
         Return the |HeaderFooterBody| object corresponding to the related
         part identified by *rId*.
         """
-        raise NotImplementedError
+        part = self.get_related_part(rId)
+        return part.body
+
+    def get_related_part(self, rId):
+        """ HACK this isn't strictly necessary
+        just adding it because it seems much easier to mock than
+        self.rels.related_parts
+        """
+        return self.rels.related_parts[rId]
 
     def save(self, path_or_stream):
         """
diff --git a/docx/parts/header.py b/docx/parts/header.py
@@ -0,0 +1,24 @@
+# encoding: utf-8
+
+"""
+|HeaderPart| and closely related objects
+"""
+
+from __future__ import (
+    absolute_import, division, print_function, unicode_literals
+)
+
+from ..header import HeaderFooterBody
+from ..opc.part import XmlPart
+
+
+class HeaderPart(XmlPart):
+    @property
+    def body(self):
+        """
+        A |HeaderFooterBody| proxy object for the `w:hdr` element in this part,
+        """
+        # TODO write CT_HeaderFooter
+        # element = CT_HeaderFooter(self.element)
+        # how to access parent here? is it necessary?
+        return HeaderFooterBody(self.element, None)
diff --git a/features/steps/header.py b/features/steps/header.py
@@ -39,7 +39,7 @@ def then_header_body_contains_the_text_of_the_header(context):
 @then('header.body is a BlockItemContainer object')
 def then_header_body_is_a_BlockItemContainer_object(context):
     header = context.header
-    assert type(header.body).__name__ == 'BlockItemContainer'
+    assert type(header.body).__name__ == 'HeaderFooterBody'
 
 
 @then('header.is_linked_to_previous is {value}')
diff --git a/tests/test_header.py b/tests/test_header.py
@@ -13,9 +13,10 @@
 from docx.enum.header import WD_HEADER_FOOTER
 from docx.header import _BaseHeaderFooter, Header, HeaderFooterBody
 from docx.parts.document import DocumentPart
+from docx.parts.header import HeaderPart
 
 from .unitutil.cxml import element
-from .unitutil.mock import call, instance_mock, property_mock
+from .unitutil.mock import call, instance_mock, method_mock, property_mock
 
 
 class Describe_BaseHeaderFooter(object):
@@ -30,8 +31,28 @@ def it_provides_access_to_its_body(self, body_fixture):
         assert header.part.related_hdrftr_body.call_args_list == calls
         assert body == expected_value
 
+    def it_provides_access_to_the_related_hdrftr_body(self, hdrftr_fixture):
+        document_part, get_related_parts, header_part_ = hdrftr_fixture
+        rId = 'rId1'
+        body = document_part.related_hdrftr_body(rId)
+        get_related_parts.assert_called_once_with(rId)
+        assert body == header_part_.body
+
     # fixtures -------------------------------------------------------
 
+    @pytest.fixture
+    def hdrftr_fixture(self, request, header_part_, body_):
+        header_part_.body = body_
+        document_part = DocumentPart(None, None, None, None)
+
+        get_related_part = method_mock(
+            request,
+            DocumentPart,
+            'get_related_part')
+        get_related_part.return_value = header_part_
+
+        return document_part, get_related_part, header_part_
+
     @pytest.fixture(params=[
         ('w:sectPr',                                             None),
         ('w:sectPr/w:headerReference{w:type=even,r:id=rId6}',    None),
@@ -64,6 +85,10 @@ def body_(self, request):
     def document_part_(self, request):
         return instance_mock(request, DocumentPart)
 
+    @pytest.fixture
+    def header_part_(self, request):
+        return instance_mock(request, HeaderPart)
+
     @pytest.fixture
     def part_prop_(self, request, document_part_):
         return property_mock(
