Read overrides from content types

2 files changed, 39 insertions, 6 deletions

diff --git a/mammoth/docx/content_types_xml.py b/mammoth/docx/content_types_xml.py
index a2da8e0..a2eb14c 100644
--- a/mammoth/docx/content_types_xml.py
+++ b/mammoth/docx/content_types_xml.py
@@ -1,9 +1,13 @@
 def read_content_types_xml_element(element):
     extension_defaults = dict(map(
         _read_default,
-        element.find_children("content-types:Default"))
-    )
-    return ContentTypes(extension_defaults)
+        element.find_children("content-types:Default")
+    ))
+    overrides = dict(map(
+        _read_override,
+        element.find_children("content-types:Override")
+    ))
+    return ContentTypes(extension_defaults, overrides)
 
 
 def _read_default(element):
@@ -12,10 +16,20 @@ def _read_default(element):
     return extension, content_type
 
 
+def _read_override(element):
+    part_name = element.attributes["PartName"]
+    content_type = element.attributes["ContentType"]
+    return part_name.lstrip("/"), content_type
+
+
 class ContentTypes(object):
-    def __init__(self, extension_defaults):
+    def __init__(self, extension_defaults, overrides):
         self._extension_defaults = extension_defaults
+        self._overrides = overrides
     
     def find_content_type(self, path):
-        extension = path.rpartition(".")[2]
-        return self._extension_defaults.get(extension)
+        if path in self._overrides:
+            return self._overrides[path]
+        else:
+            extension = path.rpartition(".")[2]
+            return self._extension_defaults.get(extension)
diff --git a/tests/docx/content_types_xml_tests.py b/tests/docx/content_types_xml_tests.py
index 56b5a86..1c26b34 100644
--- a/tests/docx/content_types_xml_tests.py
+++ b/tests/docx/content_types_xml_tests.py
@@ -17,3 +17,22 @@ def content_type_is_based_on_default_for_extension_if_there_is_no_override():
         "image/png",
         content_types.find_content_type("word/media/hat.png"),
     )
+
+
+@istest
+def content_type_is_based_on_default_for_extension_if_there_is_no_override():
+    element = xml_element("content-types:Types", {}, [
+        xml_element("content-types:Default", {
+            "Extension": "png",
+            "ContentType": "image/png",
+        }),
+        xml_element("content-types:Override", {
+            "PartName": "/word/media/hat.png",
+            "ContentType": "image/hat"
+        }),
+    ])
+    content_types = read_content_types_xml_element(element)
+    assert_equal(
+        "image/hat",
+        content_types.find_content_type("word/media/hat.png"),
+    )