diff options
| -rw-r--r-- | mammoth/docx/style_map.py | 7 | ||||
| -rw-r--r-- | mammoth/zips.py | 28 | ||||
| -rw-r--r-- | tests/docx/style_map_tests.py | 34 |
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) |