Update content types when embedding style map

1 files changed, 19 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