Read inline images

6 files changed, 135 insertions, 8 deletions

diff --git a/mammoth/documents.py b/mammoth/documents.py
index c8082d3..851755e 100644
--- a/mammoth/documents.py
+++ b/mammoth/documents.py
@@ -4,10 +4,19 @@ Document = collections.namedtuple("Document", ["children"])
 Paragraph = collections.namedtuple("Paragraph", ["children", "style_name", "numbering"])
 Run = collections.namedtuple("Run", ["children", "style_name", "is_bold", "is_italic"])
 Text = collections.namedtuple("Text", ["value"])
+
+
 class Tab(object):
     pass
 
 
+class Image(object):
+    def __init__(self, alt_text, content_type, open):
+        self.alt_text = alt_text
+        self.content_type = content_type
+        self.open = open
+
+
 document = Document
 
 def paragraph(children, style_name=None, numbering=None):
@@ -17,6 +26,11 @@ def run(children, style_name=None, is_bold=None, is_italic=None):
     return Run(children, style_name, bool(is_bold), bool(is_italic))
 
 text = Text
+
 _tab = Tab()
+
 def tab():
     return _tab
+
+
+image = Image
diff --git a/mammoth/docx/document_xml.py b/mammoth/docx/document_xml.py
index 71d744c..f8f1656 100644
--- a/mammoth/docx/document_xml.py
+++ b/mammoth/docx/document_xml.py
@@ -1,14 +1,16 @@
+import os
+
 from .. import documents
 from .. import results
 from .. import lists
 from .xmlparser import node_types
 
 
-def read_document_xml_element(element, numbering=None):
-    reader = _create_reader(numbering)
+def read_document_xml_element(element, numbering=None, content_types=None, relationships=None, docx_file=None):
+    reader = _create_reader(numbering=numbering, content_types=content_types, relationships=relationships, docx_file=docx_file)
     return reader(element)
 
-def _create_reader(numbering):
+def _create_reader(numbering, content_types, relationships, docx_file):
     _handlers = {}
     _ignored_elements = set([
         "w:bookmarkStart",
@@ -19,7 +21,9 @@ def _create_reader(numbering):
         "w:commentRangeStart",
         "w:commentRangeEnd",
         "w:commentReference",
-        "w:del"
+        "w:del",
+        "w:pPr",
+        "w:rPr",
     ])
 
     def handler(name):
@@ -96,6 +100,35 @@ def _create_reader(numbering):
     def ins(element):
         return _read_xml_elements(element.children)
     
+    
+    @handler("w:drawing")
+    def drawing(element):
+        return _read_xml_elements(element.children)
+    
+    @handler("wp:inline")
+    def inline(element):
+        alt_text = element.find_child_or_null("wp:docPr").attributes.get("descr")
+        blips = element.find_children("a:graphic") \
+            .find_children("a:graphicData") \
+            .find_children("pic:pic") \
+            .find_children("pic:blipFill") \
+            .find_children("a:blip")
+        return _read_blips(blips, alt_text)
+    
+    def _read_blips(blips, alt_text):
+        return results.combine(map(lambda blip: _read_blip(blip, alt_text), blips))
+    
+    def _read_blip(element, alt_text):
+        relationship_id = element.attributes["r:embed"]
+        image_path = os.path.join("word", relationships[relationship_id].target)
+        content_type = content_types.find_content_type(image_path)
+        
+        def open_image():
+            return docx_file.open(image_path)
+        
+        image = documents.image(alt_text=alt_text, content_type=content_type, open=open_image)
+        return results.success(image)
+    
     def read(element):
         handler = _handlers.get(element.name)
         if handler is None:
diff --git a/mammoth/docx/relationships_xml.py b/mammoth/docx/relationships_xml.py
new file mode 100644
index 0000000..75069fc
--- /dev/null
+++ b/mammoth/docx/relationships_xml.py
@@ -0,0 +1,10 @@
+class Relationships(object):
+    def __init__(self, relationships):
+        self._relationships = relationships
+    
+    def __getitem__(self, key):
+        return self._relationships[key]
+
+class Relationship(object):
+    def __init__(self, target):
+        self.target = target
diff --git a/mammoth/docx/xmlparser.py b/mammoth/docx/xmlparser.py
index 50c80db..74d8730 100644
--- a/mammoth/docx/xmlparser.py
+++ b/mammoth/docx/xmlparser.py
@@ -14,7 +14,22 @@ class XmlElement(XmlElementBase):
         
     
     def find_children(self, name):
-        return filter(lambda child: child.name == name, self.children)
+        return XmlElementList(filter(lambda child: child.name == name, self.children))
+
+
+class XmlElementList(object):
+    def __init__(self, elements):
+        self._elements = elements
+        
+    def __iter__(self):
+        return iter(self._elements)
+    
+    def find_children(self, name):
+        children = []
+        for element in self._elements:
+            for child in element.find_children(name):
+                children.append(child)
+        return XmlElementList(children)
 
 
 class NullXmlElement(object):
diff --git a/test-requirements.txt b/test-requirements.txt
index 02f2252..9258187 100644
--- a/test-requirements.txt
+++ b/test-requirements.txt
@@ -1 +1,2 @@
 nose>=1.2.1,<2.0
+funk>=0.3.1,<0.4
diff --git a/tests/docx/document_xml_tests.py b/tests/docx/document_xml_tests.py
index 69286eb..36c663b 100644
--- a/tests/docx/document_xml_tests.py
+++ b/tests/docx/document_xml_tests.py
@@ -1,9 +1,13 @@
+import io
+
 from nose.tools import istest, assert_equal
+import funk
 
 from mammoth import documents, results
 from mammoth.docx.xmlparser import element as xml_element, text as xml_text
 from mammoth.docx.document_xml import read_document_xml_element
 from mammoth.docx.numbering_xml import Numbering, NumberingLevel
+from mammoth.docx.relationships_xml import Relationships, Relationship
 
 
 @istest
@@ -124,6 +128,38 @@ class ReadXmlElementTests(object):
             documents.paragraph([documents.run([])]),
             _read_and_get_document_xml_element(element)
         )
+        
+    @istest
+    @funk.with_context
+    def can_read_inline_pictures(self, context):
+        drawing_element = _create_inline_image(
+            relationship_id="rId5",
+            description="It's a hat",
+        )
+        
+        image_bytes = b"Not an image at all!"
+        
+        relationships = Relationships({
+            "rId5": Relationship(target="media/hat.png")
+        })
+        
+        docx_file = context.mock()
+        funk.allows(docx_file).open("word/media/hat.png").returns(io.BytesIO(image_bytes))
+        
+        content_types = context.mock()
+        funk.allows(content_types).find_content_type("word/media/hat.png").returns("image/png")
+        
+        image = _read_and_get_document_xml_element(
+            drawing_element,
+            content_types=content_types,
+            relationships=relationships,
+            docx_file=docx_file,
+        )[0]
+        assert_equal(documents.Image, type(image))
+        assert_equal("It's a hat", image.alt_text)
+        assert_equal("image/png", image.content_type)
+        with image.open() as image_file:
+            assert_equal(image_bytes, image_file.read())
     
     @istest
     def ignored_elements_are_ignored_without_message(self):
@@ -142,18 +178,20 @@ class ReadXmlElementTests(object):
     @istest
     def unrecognised_elements_are_ignored(self):
         element = xml_element("w:huh", {}, [])
-        assert_equal(None, _read_and_get_document_xml_element(element))
+        assert_equal(None, read_document_xml_element(element).value)
     
     @istest
     def unrecognised_children_are_ignored(self):
         element = xml_element("w:r", {}, [_text_element("Hello!"), xml_element("w:huh", {}, [])])
         assert_equal(
             documents.run([documents.Text("Hello!")]),
-            _read_and_get_document_xml_element(element)
+            read_document_xml_element(element).value
         )
 
 def _read_and_get_document_xml_element(*args, **kwargs):
-    return read_document_xml_element(*args, **kwargs).value
+    result = read_document_xml_element(*args, **kwargs)
+    assert_equal([], result.messages)
+    return result.value
 
 
 def _document_element_with_text(text):
@@ -173,3 +211,19 @@ def _run_element_with_text(text):
 def _text_element(value):
     return xml_element("w:t", {}, [xml_text(value)])
     
+
+def _create_inline_image(description, relationship_id):
+    return xml_element("w:drawing", {}, [
+        xml_element("wp:inline", {}, [
+            xml_element("wp:docPr", {"descr": description}),
+            xml_element("a:graphic", {}, [
+                xml_element("a:graphicData", {}, [
+                    xml_element("pic:pic", {}, [
+                        xml_element("pic:blipFill", {}, [
+                            xml_element("a:blip", {"r:embed": relationship_id})
+                        ])
+                    ])
+                ])
+            ])
+        ])
+    ])