Read footnotes and store in Document

3 files changed, 56 insertions, 6 deletions

diff --git a/mammoth/documents.py b/mammoth/documents.py
index 326fbeb..ed98aaa 100644
--- a/mammoth/documents.py
+++ b/mammoth/documents.py
@@ -1,7 +1,7 @@
 import dodge
 
 
-Document = dodge.data_class("Document", ["children"])
+Document = dodge.data_class("Document", ["children", "footnotes"])
 Paragraph = dodge.data_class("Paragraph", ["children", "style_id", "style_name", "numbering"])
 Run = dodge.data_class("Run", ["children", "style_id", "style_name", "is_bold", "is_italic", "is_underline"])
 Text = dodge.data_class("Text", ["value"])
@@ -22,7 +22,10 @@ class Image(object):
         self.open = open
 
 
-document = Document
+def document(children, footnotes=None):
+    if footnotes is None:
+        footnotes = Footnotes({})
+    return Document(children, footnotes)
 
 def paragraph(children, style_id=None, style_name=None, numbering=None):
     return Paragraph(children, style_id, style_name, numbering)
@@ -49,3 +52,22 @@ def numbering_level(level_index, is_ordered):
     return _NumberingLevel(str(level_index), bool(is_ordered))
 
 _NumberingLevel = dodge.data_class("NumberingLevel", ["level_index", "is_ordered"])
+
+
+footnote = Footnote = dodge.data_class("Footnote", ["id", "body"])
+
+
+class Footnotes(object):
+    def __init__(self, footnotes):
+        self._footnotes = footnotes
+    
+    def find_footnote_by_id(self, id):
+        return self._footnotes[id]
+    
+    def __eq__(self, other):
+        return isinstance(other, Footnotes) and self._footnotes == other._footnotes
+
+    def __ne__(self, other):
+        return not (self == other)
+
+footnotes = Footnotes
diff --git a/mammoth/docx/document_xml.py b/mammoth/docx/document_xml.py
index f96098e..68b4bb3 100644
--- a/mammoth/docx/document_xml.py
+++ b/mammoth/docx/document_xml.py
@@ -7,17 +7,28 @@ from .. import lists
 from .xmlparser import node_types
 
 
-def read_document_xml_element(element, numbering=None, content_types=None, relationships=None, styles=None, docx_file=None):
+def read_document_xml_element(element,
+        numbering=None,
+        content_types=None,
+        relationships=None,
+        styles=None,
+        footnote_elements=None,
+        docx_file=None):
+    
+    if footnote_elements is None:
+        footnote_elements = []
+    
     reader = _create_reader(
         numbering=numbering,
         content_types=content_types,
         relationships=relationships,
         styles=styles,
+        footnote_elements=footnote_elements,
         docx_file=docx_file,
     )
     return reader(element)
 
-def _create_reader(numbering, content_types, relationships, styles, docx_file):
+def _create_reader(numbering, content_types, relationships, styles, footnote_elements, docx_file):
     _handlers = {}
     _ignored_elements = set([
         "w:bookmarkStart",
@@ -115,8 +126,21 @@ def _create_reader(numbering, content_types, relationships, styles, docx_file):
     @handler("w:document")
     def document(element):
         body_element = _find_child(element, "w:body")
-        return _read_xml_elements(body_element.children) \
-            .map(documents.document)
+        children_result = _read_xml_elements(body_element.children)
+        footnotes_result = results.combine(map(_read_footnote, footnote_elements))
+        
+        return results.map(
+            lambda children, footnotes: documents.document(
+                children,
+                documents.footnotes(dict((footnote.id, footnote) for footnote in footnotes))
+            ),
+            children_result, footnotes_result
+        )
+    
+    
+    def _read_footnote(element):
+        return _read_xml_elements(element.body).map(
+            lambda body: documents.footnote(element.id, body))
     
     
     @handler("w:tab")
diff --git a/mammoth/results.py b/mammoth/results.py
index b4bedbd..9112296 100644
--- a/mammoth/results.py
+++ b/mammoth/results.py
@@ -34,3 +34,7 @@ def combine(results):
             messages.append(message)
         
     return Result(values, messages)
+
+
+def map(func, *args):
+    return combine(args).map(lambda values: func(*values))