Update content types when embedding style map

2 files changed, 46 insertions, 3 deletions

diff --git a/mammoth/docx/style_map.py b/mammoth/docx/style_map.py
index 1d35a63..9d2646a 100644
--- a/mammoth/docx/style_map.py
+++ b/mammoth/docx/style_map.py
@@ -6,15 +6,18 @@ from ..zips import open_zip, update_zip
 _style_map_path = "mammoth/style-map"
 _style_map_absolute_path = "/" + _style_map_path
 _relationships_path = "word/_rels/document.xml.rels"
+_content_types_path = "[Content_Types].xml"
 
 
 def write_style_map(fileobj, style_map):
     with open_zip(fileobj, "r") as zip_file:
-        relationships_xml = _generate_relationships_xml(zip_file.read_str("word/_rels/document.xml.rels"))
+        relationships_xml = _generate_relationships_xml(zip_file.read_str(_relationships_path))
+        content_types_xml = _generate_content_types_xml(zip_file.read_str(_content_types_path))
     
     update_zip(fileobj, {
         _style_map_path: style_map,
-        _relationships_path: relationships_xml
+        _relationships_path: relationships_xml,
+        _content_types_path: content_types_xml,
     })
 
 def _generate_relationships_xml(relationships_xml):
@@ -32,6 +35,19 @@ def _generate_relationships_xml(relationships_xml):
     return ElementTree.tostring(relationships, "UTF-8")
 
 
+def _generate_content_types_xml(content_types_xml):
+    content_types_uri = "http://schemas.openxmlformats.org/package/2006/content-types"
+    override_name = "{" + content_types_uri + "}Override"
+    
+    types = ElementTree.fromstring(content_types_xml)
+    _add_or_update_element(types, override_name, "PartName", {
+        "PartName": _style_map_absolute_path,
+        "ContentType": "text/prs.mammoth.style-map",
+    })
+    
+    return ElementTree.tostring(types, "UTF-8")
+
+
 def _add_or_update_element(parent, name, identifying_attribute, attributes):
     existing_child = _find_child(parent, name, identifying_attribute, attributes)
     if existing_child is None:
@@ -41,7 +57,7 @@ def _add_or_update_element(parent, name, identifying_attribute, attributes):
     
 
 def _find_child(parent, name, identifying_attribute, attributes):
-    for element in parent.iter():
+    for element in parent.getiterator():
         if element.tag == name and element.get(identifying_attribute) == attributes.get(identifying_attribute):
             return element
 
diff --git a/tests/docx/style_map_tests.py b/tests/docx/style_map_tests.py
index 6c9e97b..b774b66 100644
--- a/tests/docx/style_map_tests.py
+++ b/tests/docx/style_map_tests.py
@@ -41,6 +41,12 @@ def embedded_style_map_is_referenced_in_relationships():
     fileobj = _normal_docx()
     write_style_map(fileobj, "p => h1")
     assert_equal(expected_relationships_xml, _read_relationships_xml(fileobj))
+    
+@istest
+def embedded_style_map_has_override_content_type_in_content_types_xml():
+    fileobj = _normal_docx()
+    write_style_map(fileobj, "p => h1")
+    assert_equal(expected_content_types_xml, _read_content_types_xml(fileobj))
 
 
 @istest
@@ -52,6 +58,7 @@ def can_overwrite_existing_style_map():
         assert_equal("p => h2", read_style_map(fileobj))
         _assert_no_duplicates(zip_file._zip_file.namelist())
         assert_equal(expected_relationships_xml, _read_relationships_xml(fileobj))
+        assert_equal(expected_content_types_xml, _read_content_types_xml(fileobj))
 
 
 def _read_relationships_xml(fileobj):
@@ -62,6 +69,14 @@ def _read_relationships_xml(fileobj):
         )
 
 
+def _read_content_types_xml(fileobj):
+    with open_zip(fileobj, "r") as zip_file:
+        return xml.parse_xml(
+            io.StringIO(zip_file.read_str("[Content_Types].xml")),
+            [("ct", "http://schemas.openxmlformats.org/package/2006/content-types")],
+        )
+
+
 original_relationships_xml = ('<?xml version="1.0" encoding="UTF-8" standalone="yes"?>' +
     '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">' +
     '<Relationship Id="rId3" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/settings" Target="settings.xml"/>' +
@@ -72,6 +87,17 @@ expected_relationships_xml = xml.element("r:Relationships", {}, [
     xml.element("r:Relationship", {"Id": "rMammothStyleMap", "Type": "http://schemas.zwobble.org/mammoth/style-map", "Target": "/mammoth/style-map"}),
 ])
 
+original_content_types_xml = ('<?xml version="1.0" encoding="UTF-8" standalone="yes"?>' +
+    '<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">' +
+    '<Default Extension="png" ContentType="image/png"/>' +
+    '</Types>'
+)
+
+expected_content_types_xml = xml.element("ct:Types", {}, [
+    xml.element("ct:Default", {"Extension": "png", "ContentType": "image/png"}),
+    xml.element("ct:Override", {"PartName": "/mammoth/style-map", "ContentType": "text/prs.mammoth.style-map"}),
+])
+
 
 def _normal_docx():
     fileobj = io.BytesIO()
@@ -79,6 +105,7 @@ def _normal_docx():
     try:
         zip_file.writestr("placeholder", "placeholder")
         zip_file.writestr("word/_rels/document.xml.rels", original_relationships_xml)
+        zip_file.writestr("[Content_Types].xml", original_content_types_xml)
         expected_relationships_xml
     finally:
         zip_file.close()