Read footnotes and store in Document

4 files changed, 71 insertions, 7 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))
diff --git a/tests/docx/document_xml_tests.py b/tests/docx/document_xml_tests.py
index 963a75f..b5b1e19 100644
--- a/tests/docx/document_xml_tests.py
+++ b/tests/docx/document_xml_tests.py
@@ -9,6 +9,7 @@ from mammoth.docx.document_xml import read_document_xml_element
 from mammoth.docx.numbering_xml import Numbering
 from mammoth.docx.relationships_xml import Relationships, Relationship
 from mammoth.docx.styles_xml import Styles, Style
+from mammoth.docx.footnotes_xml import FootnoteElement
 
 
 @istest
@@ -38,7 +39,7 @@ class ReadXmlElementTests(object):
     def can_read_text_within_document(self):
         element = _document_element_with_text("Hello!")
         assert_equal(
-            documents.Document([documents.paragraph([documents.run([documents.Text("Hello!")])])]),
+            documents.document([documents.paragraph([documents.run([documents.Text("Hello!")])])]),
             _read_and_get_document_xml_element(element)
         )
         
@@ -300,6 +301,19 @@ class ReadXmlElementTests(object):
         assert_equal("image/png", image.content_type)
         with image.open() as image_file:
             assert_equal(image_bytes, image_file.read())
+        
+    @istest
+    def footnotes_of_document_are_read(self):
+        paragraph_xml = xml_element("w:p")
+        footnotes = [FootnoteElement("4", [paragraph_xml])]
+        
+        body_xml = xml_element("w:body")
+        document_xml = xml_element("w:document", {}, [body_xml])
+        
+        document = _read_and_get_document_xml_element(document_xml, footnote_elements=footnotes)
+        footnote = document.footnotes.find_footnote_by_id("4")
+        assert_equal("4", footnote.id)
+        assert isinstance(footnote.body[0], documents.Paragraph)
     
     @istest
     def ignored_elements_are_ignored_without_message(self):