Change footnote reader to generic note reader

1 files changed, 19 insertions, 13 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")