Support overwriting embedded style map

2 files changed, 29 insertions, 6 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)