Change footnote reader to generic note reader

2 files changed, 21 insertions, 15 deletions

diff --git a/mammoth/docx/footnotes_xml.py b/mammoth/docx/footnotes_xml.py
index 1c596ab..c477553 100644
--- a/mammoth/docx/footnotes_xml.py
+++ b/mammoth/docx/footnotes_xml.py
@@ -1,23 +1,29 @@
 from .. import lists
 
 
-def read_footnotes_xml_element(element):
-    footnote_elements = lists.filter(
-        _is_footnote_element,
-        element.find_children("w:footnote"),
-    )
-    return lists.map(_read_footnote_element, footnote_elements)
+def create_reader(note_type):
+    def read_notes_xml_element(element):
+        note_elements = lists.filter(
+            _is_note_element,
+            element.find_children("w:" + note_type),
+        )
+        return lists.map(_read_note_element, note_elements)
 
 
-def _is_footnote_element(element):
-    return element.attributes.get("w:type") not in ["continuationSeparator", "separator"]
+    def _is_note_element(element):
+        return element.attributes.get("w:type") not in ["continuationSeparator", "separator"]
 
 
-def _read_footnote_element(element):
-    return FootnoteElement(element.attributes["w:id"], element.children)
+    def _read_note_element(element):
+        return NoteElement(note_type, element.attributes["w:id"], element.children)
+    
+    return read_notes_xml_element
 
-
-class FootnoteElement(object):
-    def __init__(self, id, body):
+class NoteElement(object):
+    def __init__(self, note_type, id, body):
+        self.note_type = note_type
         self.id = id
         self.body = body
+
+
+read_footnotes_xml_element = create_reader("footnote")
diff --git a/tests/docx/document_xml_tests.py b/tests/docx/document_xml_tests.py
index 4c1f330..2cb13c7 100644
--- a/tests/docx/document_xml_tests.py
+++ b/tests/docx/document_xml_tests.py
@@ -9,7 +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
+from mammoth.docx.footnotes_xml import NoteElement
 
 
 @istest
@@ -315,7 +315,7 @@ class ReadXmlElementTests(object):
     @istest
     def footnotes_of_document_are_read(self):
         paragraph_xml = xml_element("w:p")
-        footnotes = [FootnoteElement("4", [paragraph_xml])]
+        footnotes = [NoteElement("footnote", "4", [paragraph_xml])]
         
         body_xml = xml_element("w:body")
         document_xml = xml_element("w:document", {}, [body_xml])