Support overwriting embedded style map

3 files changed, 61 insertions, 8 deletions

diff --git a/mammoth/docx/style_map.py b/mammoth/docx/style_map.py
index 2841830..94c32af 100644
--- a/mammoth/docx/style_map.py
+++ b/mammoth/docx/style_map.py
@@ -1,12 +1,13 @@
-from ..zips import open_zip
+from ..zips import open_zip, update_zip
 
 
 _style_map_path = "mammoth/style-map"
 
 
 def write_style_map(fileobj, style_map):
-    with open_zip(fileobj, "a") as zip_file:
-        zip_file.write_str(_style_map_path, style_map)
+    update_zip(fileobj, {
+        _style_map_path: style_map
+    })
 
 
 def read_style_map(fileobj):
diff --git a/mammoth/zips.py b/mammoth/zips.py
index 24c6b1c..432be5d 100644
--- a/mammoth/zips.py
+++ b/mammoth/zips.py
@@ -1,4 +1,6 @@
 import contextlib
+import io
+import shutil
 
 from zipfile import ZipFile
 
@@ -27,8 +29,28 @@ class _Zip(object):
         except KeyError:
             return False
 
-    def write_str(self, name, value):
-        self._zip_file.writestr(name, value)
-    
     def read_str(self, name):
         return self._zip_file.read(name).decode("utf8")
+
+
+def update_zip(fileobj, files):
+    source = ZipFile(fileobj, "r")
+    try:
+        destination_fileobj = io.BytesIO()
+        destination = ZipFile(destination_fileobj, "w")
+        try:
+            names = set(source.namelist()) | set(files.keys())
+            for name in names:
+                if name in files:
+                    contents = files[name]
+                else:
+                    contents = source.read(name)
+                destination.writestr(name, contents)
+        finally:
+            destination.close()
+    finally:
+        source.close()
+    
+    fileobj.seek(0)
+    destination_fileobj.seek(0)
+    shutil.copyfileobj(destination_fileobj, fileobj)
diff --git a/tests/docx/style_map_tests.py b/tests/docx/style_map_tests.py
index f298e1d..5cf2292 100644
--- a/tests/docx/style_map_tests.py
+++ b/tests/docx/style_map_tests.py
@@ -1,4 +1,5 @@
 import io
+from zipfile import ZipFile
 
 from nose.tools import istest, assert_equal
 
@@ -13,6 +14,13 @@ def reading_embedded_style_map_on_document_without_embedded_style_map_returns_no
 
 
 @istest
+def writing_style_map_preserves_unrelated_files():
+    fileobj = _normal_docx()
+    write_style_map(fileobj, "p => h1")
+    with open_zip(fileobj, "r") as zip_file:
+        assert_equal("placeholder", zip_file.read_str("placeholder"))
+
+@istest
 def embedded_style_map_can_be_read_after_being_written():
     fileobj = _normal_docx()
     write_style_map(fileobj, "p => h1")
@@ -27,8 +35,30 @@ def embedded_style_map_is_written_to_separate_file():
         assert_equal("p => h1", zip_file.read_str("mammoth/style-map"))
 
 
+@istest
+def can_overwrite_existing_style_map():
+    fileobj = _normal_docx()
+    write_style_map(fileobj, "p => h1")
+    write_style_map(fileobj, "p => h2")
+    with open_zip(fileobj, "r") as zip_file:
+        assert_equal("p => h2", read_style_map(fileobj))
+        _assert_no_duplicates(zip_file._zip_file.namelist())
+
+
 def _normal_docx():
     fileobj = io.BytesIO()
-    with open_zip(fileobj, "w") as zip_file:
-        zip_file.write_str("placeholder", "placeholder")
+    zip_file = ZipFile(fileobj, "w")
+    try:
+        zip_file.writestr("placeholder", "placeholder")
+    finally:
+        zip_file.close()
     return fileobj
+
+
+def _assert_no_duplicates(values):
+    counts = {}
+    for value in values:
+        counts[value] = counts.get(value, 0) + 1
+    for value, count in counts.items():
+        if count != 1:
+            assert False, "{0} has count of {1}".format(value, count)