Handle relationships in notes

1 files changed, 26 insertions, 14 deletions

diff --git a/mammoth/docx/__init__.py b/mammoth/docx/__init__.py
index 3e7c0e0..ec8d363 100644
--- a/mammoth/docx/__init__.py
+++ b/mammoth/docx/__init__.py
@@ -1,13 +1,14 @@
 import zipfile
 import contextlib
 
+from .. import results, lists
 from .xmlparser import parse_xml
 from .document_xml import read_document_xml_element
 from .content_types_xml import read_content_types_xml_element
-from .relationships_xml import read_relationships_xml_element
+from .relationships_xml import read_relationships_xml_element, Relationships
 from .numbering_xml import read_numbering_xml_element, Numbering
 from .styles_xml import read_styles_xml_element
-from .notes_xml import read_footnotes_xml_element, read_endnotes_xml_element
+from .notes_xml import create_footnotes_reader, create_endnotes_reader
 from . import body_xml
 
 
@@ -25,23 +26,32 @@ _namespaces = [
 
 def read(fileobj):
     zip_file = zipfile.ZipFile(fileobj)
-    
     body_readers = _body_readers(zip_file)
-    body_reader = body_readers("document")
     
-    footnote_elements = _try_read_entry_or_default(
-        zip_file, "word/footnotes.xml", read_footnotes_xml_element, default=[])
+    return _read_notes(zip_file, body_readers).bind(lambda notes:
+        _read_document(zip_file, body_readers, notes))
+
+
+def _read_notes(zip_file, body_readers):
+    empty_result = results.success([])
+    
+    read_footnotes_xml = create_footnotes_reader(body_readers("footnotes"))
+    footnotes = _try_read_entry_or_default(
+        zip_file, "word/footnotes.xml", read_footnotes_xml, default=empty_result)
+    
+    read_endnotes_xml = create_endnotes_reader(body_readers("endnotes"))
+    endnotes = _try_read_entry_or_default(
+        zip_file, "word/endnotes.xml", read_endnotes_xml, default=empty_result)
     
-    endnote_elements = _try_read_entry_or_default(
-        zip_file, "word/endnotes.xml", read_endnotes_xml_element, default=[])
+    return results.combine([footnotes, endnotes]).map(lists.collect)
     
+def _read_document(zip_file, body_readers, notes):
     with _open_entry(zip_file, "word/document.xml") as document_fileobj:
         document_xml = _parse_docx_xml(document_fileobj)
         return read_document_xml_element(
             document_xml,
-            body_reader=body_reader,
-            footnote_elements=footnote_elements,
-            endnote_elements=endnote_elements,
+            body_reader=body_readers("document"),
+            notes=notes,
         )
 
 
@@ -56,9 +66,11 @@ def _body_readers(zip_file):
         styles = read_styles_xml_element(_parse_docx_xml(styles_fileobj))
     
     def for_name(name):
-        with _open_entry(zip_file, "word/_rels/{0}.xml.rels".format(name)) as relationships_fileobj:
-            relationships = read_relationships_xml_element(_parse_docx_xml(relationships_fileobj))
-        
+        relationships_path = "word/_rels/{0}.xml.rels".format(name)
+        relationships = _try_read_entry_or_default(
+            zip_file, relationships_path, read_relationships_xml_element,
+            default=Relationships({}))
+            
         return body_xml.reader(
             numbering=numbering,
             content_types=content_types,