Handle relationships in notes

1 files changed, 5 insertions, 11 deletions

diff --git a/mammoth/docx/document_xml.py b/mammoth/docx/document_xml.py
index 3cdbec2..1c15001 100644
--- a/mammoth/docx/document_xml.py
+++ b/mammoth/docx/document_xml.py
@@ -1,19 +1,13 @@
 from .. import documents
-from .. import results
 
 
 def read_document_xml_element(element,
         body_reader,
-        footnote_elements=None,
-        endnote_elements=None):
+        notes=None):
     
-    note_elements = (footnote_elements or []) + (endnote_elements or [])
+    if notes is None:
+        notes = []
     
-    def _read_note(element):
-        return body_reader.read_all(element.body).map(
-            lambda body: documents.note(element.note_type, element.id, body))
-
     body_element = element.find_child("w:body")
-    children_result = body_reader.read_all(body_element.children)
-    notes_result = results.combine(map(_read_note, note_elements)).map(documents.notes)
-    return results.map(documents.document, children_result, notes_result)
+    return body_reader.read_all(body_element.children) \
+        .map(lambda children: documents.document(children, documents.notes(notes)))