Update relationships when embedded style map

2 files changed, 67 insertions, 1 deletions

diff --git a/mammoth/docx/style_map.py b/mammoth/docx/style_map.py
index 94c32af..1d35a63 100644
--- a/mammoth/docx/style_map.py
+++ b/mammoth/docx/style_map.py
@@ -1,14 +1,50 @@
+from xml.etree import ElementTree
+
 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"
 
 
 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"))
+    
     update_zip(fileobj, {
-        _style_map_path: style_map
+        _style_map_path: style_map,
+        _relationships_path: relationships_xml
+    })
+
+def _generate_relationships_xml(relationships_xml):
+    schema = "http://schemas.zwobble.org/mammoth/style-map";
+    relationships_uri = "http://schemas.openxmlformats.org/package/2006/relationships"
+    relationship_element_name = "{" + relationships_uri + "}Relationship"
+    
+    relationships = ElementTree.fromstring(relationships_xml)
+    _add_or_update_element(relationships, relationship_element_name, "Id", {
+        "Id": "rMammothStyleMap",
+        "Type": schema,
+        "Target": _style_map_absolute_path,
     })
 
+    return ElementTree.tostring(relationships, "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:
+        ElementTree.SubElement(parent, name, attributes)
+    else:
+        existing_child.attrib = attributes
+    
+
+def _find_child(parent, name, identifying_attribute, attributes):
+    for element in parent.iter():
+        if element.tag == name and element.get(identifying_attribute) == attributes.get(identifying_attribute):
+            return element
+
 
 def read_style_map(fileobj):
     with open_zip(fileobj, "r") as zip_file:
diff --git a/tests/docx/style_map_tests.py b/tests/docx/style_map_tests.py
index 5cf2292..6c9e97b 100644
--- a/tests/docx/style_map_tests.py
+++ b/tests/docx/style_map_tests.py
@@ -5,6 +5,7 @@ from nose.tools import istest, assert_equal
 
 from mammoth.docx.style_map import write_style_map, read_style_map
 from mammoth.zips import open_zip
+from mammoth.docx import xmlparser as xml
 
 
 @istest
@@ -36,6 +37,13 @@ def embedded_style_map_is_written_to_separate_file():
 
 
 @istest
+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 can_overwrite_existing_style_map():
     fileobj = _normal_docx()
     write_style_map(fileobj, "p => h1")
@@ -43,6 +51,26 @@ def can_overwrite_existing_style_map():
     with open_zip(fileobj, "r") as zip_file:
         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))
+
+
+def _read_relationships_xml(fileobj):
+    with open_zip(fileobj, "r") as zip_file:
+        return xml.parse_xml(
+            io.StringIO(zip_file.read_str("word/_rels/document.xml.rels")),
+            [("r", "http://schemas.openxmlformats.org/package/2006/relationships")],
+        )
+
+
+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"/>' +
+    '</Relationships>')
+
+expected_relationships_xml = xml.element("r:Relationships", {}, [
+    xml.element("r:Relationship", {"Id": "rId3", "Type": "http://schemas.openxmlformats.org/officeDocument/2006/relationships/settings", "Target": "settings.xml"}),
+    xml.element("r:Relationship", {"Id": "rMammothStyleMap", "Type": "http://schemas.zwobble.org/mammoth/style-map", "Target": "/mammoth/style-map"}),
+])
 
 
 def _normal_docx():
@@ -50,6 +78,8 @@ def _normal_docx():
     zip_file = ZipFile(fileobj, "w")
     try:
         zip_file.writestr("placeholder", "placeholder")
+        zip_file.writestr("word/_rels/document.xml.rels", original_relationships_xml)
+        expected_relationships_xml
     finally:
         zip_file.close()
     return fileobj