Added support for header/footer for merge_templates #17

Refactored some tests to use the generic document extraction method Added support for adding new XML files in the docx. (new headers/footers) Added support for various generation of IDs in the docx Refactored the parts so that all docx files can be handled in a generic way. Refactored the settings using the parts.
iulica · Nov 10, 2023 · 4ddc603 · 4ddc603
1 parent 22dc108
commit 4ddc603
Show file tree

Hide file tree

Showing 11 changed files with 360 additions and 81 deletions.
diff --git a/mailmerge.py b/mailmerge.py
diff --git a/tests/test_footer.docx b/tests/test_footer.docx
diff --git a/tests/test_footnote_header_footer.docx b/tests/test_footnote_header_footer.docx
diff --git a/tests/test_footnote_header_footer.py b/tests/test_footnote_header_footer.py
@@ -1,31 +1,69 @@
 import unittest
-import tempfile
-import warnings
-import datetime
-import locale
-from os import path
-from lxml import etree
 
 from mailmerge import NAMESPACES
-from tests.utils import EtreeMixin, get_document_body_part
+from tests.utils import EtreeMixin, get_document_body_part, TEXTS_XPATH, get_document_body_parts
 
 FOOTNOTE_XPATH = "//w:footnote[@w:id = '1']/w:p/w:r/w:t/text()"
 
 class FootnoteHeaderFooterTest(EtreeMixin, unittest.TestCase):
-
+    # @TODO test missing values
+    # @TODO test if separator isn't section
+    # @TODO test headers/footers with relations
     @unittest.expectedFailure
     def test_all(self):
         values = ["one", "two", "three"]
         # header/footer/footnotes don't work with multiple replacements, only with merge
         # fix this when it is implemented
         document, root_elem = self.merge_templates(
             'test_footnote_header_footer.docx',
-            [{"fieldname": value, "footerfield": "f_" + value, "headerfield": "h_" + value}
+            [
+                {
+                "fieldname": value,
+                "footerfield": "f_" + value,
+                "headerfield": "h_" + value,
+                "footerfirst": "ff_" + value,
+                "headerfirst": "hf_" + value,
+                "footereven": "fe_" + value,
+                "headereven": "he_" + value,
+                }
                 for value in values
                 ],
-            # output="tests/test_output_footnote_header_footer.docx"
+            separator="nextPage_section",
+            # output="tests/output/test_output_footnote_header_footer.docx"
             )
 
+        footers = sorted([
+            "".join(footer_doc_tree.getroot().xpath(TEXTS_XPATH, namespaces=NAMESPACES))
+            for footer_doc_tree in get_document_body_parts(document, endswith="ftr")
+        ])
+        self.assertListEqual(footers, sorted([
+            footer
+            for value in values
+            for footer in [
+            'Footer on even page fe_%s' % value, 
+            'Footer on every page f_%s' % value,
+            'Footer on first page ff_%s' % value]] + [
+            'Footer on even page ', 
+            'Footer on every page ',
+            'Footer on first page ']
+            ))
+
+        headers = sorted([
+            "".join(header_doc_tree.getroot().xpath(TEXTS_XPATH, namespaces=NAMESPACES))
+            for header_doc_tree in get_document_body_parts(document, endswith="hdr")
+        ])
+        self.assertListEqual(headers, sorted([
+            header
+            for value in values
+            for header in [
+            'Header even: he_%s' % value,
+            'Header on every page: h_%s' % value,
+            'Header on first page: hf_%s' % value]] + [
+            'Header even: ', 
+            'Header on every page: ',
+            'Header on first page: ']
+            ))
+
         footnote_root_elem = get_document_body_part(document, "footnotes").getroot()
         footnote = "".join(footnote_root_elem.xpath(FOOTNOTE_XPATH, namespaces=NAMESPACES))
         correct_footnote = " Merge : one "
@@ -35,12 +73,57 @@ def test_only_merge(self):
         values = ["one", "two", "three"]
         document, root_elem = self.merge(
             'test_footnote_header_footer.docx',
-            [{"fieldname": value, "footerfield": "f_" + value, "headerfield": "h_" + value}
+            [
+                {
+                "fieldname": value,
+                "footerfield": "f_" + value,
+                "headerfield": "h_" + value,
+                "footerfirst": "ff_" + value,
+                "headerfirst": "hf_" + value,
+                "footereven": "fe_" + value,
+                "headereven": "he_" + value,
+                }
                 for value in values
                 ][0],
-            # output="tests/test_output_one_footnote_header_footer.docx"
+            # output="tests/output/test_output_one_footnote_header_footer.docx"
             )
 
         footnote_root_elem = get_document_body_part(document, "footnotes").getroot()
         footnote = "".join(footnote_root_elem.xpath(FOOTNOTE_XPATH, namespaces=NAMESPACES))
         self.assertEqual(footnote, " Merge : one")
+
+        footers = sorted([
+            "".join(footer_doc_tree.getroot().xpath(TEXTS_XPATH, namespaces=NAMESPACES))
+            for footer_doc_tree in get_document_body_parts(document, endswith="ftr")
+        ])
+        value = values[0]
+        self.assertListEqual(footers, [
+            'Footer on even page fe_%s' % value, 
+            'Footer on every page f_%s' % value,
+            'Footer on first page ff_%s' % value])
+
+        headers = sorted([
+            "".join(header_doc_tree.getroot().xpath(TEXTS_XPATH, namespaces=NAMESPACES))
+            for header_doc_tree in get_document_body_parts(document, endswith="hdr")
+        ])
+        self.assertListEqual(headers, [
+            'Header even: he_%s' % value,
+            'Header on every page: h_%s' % value,
+            'Header on first page: hf_%s' % value])
+
+
+    def test_footer(self):
+        values = ["one", "two"]
+        # header/footer/footnotes don't work with multiple replacements, only with merge
+        # fix this when it is implemented
+        document, root_elem = self.merge_templates(
+            'test_footer.docx',
+            [
+                {
+                "footer": value,
+                }
+                for value in values
+                ],
+            separator="nextPage_section",
+            # output="tests/output/test_footer.docx"
+            )
diff --git a/tests/test_macword2011.py b/tests/test_macword2011.py
@@ -4,7 +4,7 @@
 from lxml import etree
 
 from mailmerge import MailMerge
-from tests.utils import EtreeMixin
+from tests.utils import EtreeMixin, get_document_body_part
 
 
 class MacWord2011Test(EtreeMixin, unittest.TestCase):
@@ -57,4 +57,4 @@ def test(self):
             '/><w:pgMar w:bottom="1417" w:footer="708" w:gutter="0" w:header="708" w:left="1417" w:right="1417" '
             'w:top="1417" /><w:cols w:space="708" /><w:docGrid w:linePitch="360" /></w:sectPr></w:body></w:document>')
 
-        self.assert_equal_tree(expected_tree, list(document.parts.values())[0].getroot())
+        self.assert_equal_tree(expected_tree, get_document_body_part(document).getroot())
diff --git a/tests/test_merge_table_multipart.py b/tests/test_merge_table_multipart.py
@@ -4,7 +4,7 @@
 from lxml import etree
 
 from mailmerge import MailMerge
-from tests.utils import EtreeMixin
+from tests.utils import EtreeMixin, get_document_body_part
 
 
 class MergeTableRowsMultipartTest(EtreeMixin, unittest.TestCase):
@@ -32,10 +32,7 @@ def test_merge_rows_on_multipart_file(self):
         with tempfile.TemporaryFile() as outfile:
             self.document.write(outfile)
 
-        for part in self.document.parts.values():
-            # only check the document part 
-            if (part.getroot().tag == '{http://schemas.openxmlformats.org/wordprocessingml/2006/main}document'):
-                self.assert_equal_tree(self.expected_tree, part.getroot())
+        self.assert_equal_tree(self.expected_tree, get_document_body_part(self.document).getroot())
 
     def test_merge_unified_on_multipart_file(self):
         self.document.merge(
@@ -52,10 +49,7 @@ def test_merge_unified_on_multipart_file(self):
         with tempfile.TemporaryFile() as outfile:
             self.document.write(outfile)
 
-        for part in self.document.parts.values():
-            # only check the document part 
-            if (part.getroot().tag == '{http://schemas.openxmlformats.org/wordprocessingml/2006/main}document'):
-                self.assert_equal_tree(self.expected_tree, part.getroot())
+        self.assert_equal_tree(self.expected_tree, get_document_body_part(self.document).getroot())
 
     def tearDown(self):
         self.document.close()
diff --git a/tests/test_merge_table_rows.py b/tests/test_merge_table_rows.py
diff --git a/tests/test_multiple_elements.py b/tests/test_multiple_elements.py
@@ -4,7 +4,7 @@
 from lxml import etree
 
 from mailmerge import MailMerge
-from tests.utils import EtreeMixin
+from tests.utils import EtreeMixin, get_document_body_part
 
 
 class MultipleElementsTest(EtreeMixin, unittest.TestCase):
@@ -49,4 +49,4 @@ def test(self):
                 'w:bottom="1440" w:left="1440" w:header="708" w:footer="708" w:gutter="0"/><w:cols '
                 'w:space="708"/><w:docGrid w:linePitch="360"/></w:sectPr></w:body></w:document>')
 
-            self.assert_equal_tree(expected_tree, list(document.parts.values())[0].getroot())
+            self.assert_equal_tree(expected_tree, get_document_body_part(document).getroot())
diff --git a/tests/test_unique_id.py b/tests/test_unique_id.py
@@ -7,7 +7,7 @@ class UniqueIdsManagerTest(unittest.TestCase):
     Testing UniqueIdsManager class
     """
 
-    def test_unique_id_manager(self):
+    def test_unique_id_manager_register_id(self):
         """
         Tests if the next record field works
         """
@@ -22,3 +22,5 @@ def test_unique_id_manager(self):
 
         for type_id, obj_id, new_id in tests:
             self.assertEqual(id_man.register_id(type_id, obj_id=obj_id), new_id)
+
+        self.assertEqual(id_man.register_id_str("footer2"), "footer3")
diff --git a/tests/test_winword2010.py b/tests/test_winword2010.py
@@ -4,7 +4,7 @@
 from lxml import etree
 
 from mailmerge import MailMerge, NAMESPACES
-from tests.utils import EtreeMixin
+from tests.utils import EtreeMixin, get_document_body_part
 
 
 class Windword2010Test(EtreeMixin, unittest.TestCase):
@@ -178,5 +178,5 @@ def test(self):
             '</w:document>'  # noqa
         )
 
-        self.assert_equal_tree(expected_tree, list(document.parts.values())[0].getroot())
+        self.assert_equal_tree(expected_tree, get_document_body_part(document).getroot())
         self.assertIsNone(document.get_settings().getroot().find('{%(w)s}mailMerge' % NAMESPACES))
diff --git a/tests/utils.py b/tests/utils.py
@@ -100,7 +100,18 @@ def get_new_docx(self, replacement_parts):
 
 def get_document_body_part(document, endswith="document"):
     for part in document.parts.values():
-        if part.getroot().tag.endswith('}%s' % endswith):
-            return part
+        if part['part'].getroot().tag.endswith('}%s' % endswith):
+            return part['part']
 
     raise AssertionError("main document body not found in document.parts")
+
+def get_document_body_parts(document, endswith="document"):
+    parts = []
+    for part in document.parts.values():
+        if part['part'].getroot().tag.endswith('}%s' % endswith):
+            parts.append(part['part'])
+    for _, _, part in document.new_parts:
+        if part.getroot().tag.endswith('}%s' % endswith):
+            parts.append(part)
+
+    return parts