diff options
| -rw-r--r-- | mammoth/__init__.py | 4 | ||||
| -rw-r--r-- | mammoth/conversion.py | 5 | ||||
| -rw-r--r-- | mammoth/docx/document_xml.py | 7 | ||||
| -rw-r--r-- | scripts/mammoth | 51 | ||||
| -rw-r--r-- | tests/cli_tests.py | 12 |
5 files changed, 65 insertions, 14 deletions
diff --git a/mammoth/__init__.py b/mammoth/__init__.py index 7491fb6..65b8821 100644 --- a/mammoth/__init__.py +++ b/mammoth/__init__.py @@ -2,12 +2,12 @@ from .results import Result from . import docx, conversion, style_reader, lists -def convert_to_html(fileobj, styles=None): +def convert_to_html(fileobj, styles=None, convert_image=None): if styles is None: styles = _default_styles return docx.read(fileobj).bind(lambda document: - conversion.convert_document_element_to_html(document, styles=_read_styles(styles)) + conversion.convert_document_element_to_html(document, styles=_read_styles(styles), convert_image=convert_image) ) diff --git a/mammoth/conversion.py b/mammoth/conversion.py index b324c16..835cb6b 100644 --- a/mammoth/conversion.py +++ b/mammoth/conversion.py @@ -73,11 +73,8 @@ class DocumentConverter(object): def _convert_image(self, image, html_generator): - image_bytes = image.open() - try: + with image.open() as image_bytes: encoded_src = base64.b64encode(image_bytes.read()).decode("ascii") - finally: - image_bytes.close() attributes = { "src": "data:{0};base64,{1}".format(image.content_type, encoded_src) diff --git a/mammoth/docx/document_xml.py b/mammoth/docx/document_xml.py index fad4f87..14d36a6 100644 --- a/mammoth/docx/document_xml.py +++ b/mammoth/docx/document_xml.py @@ -1,4 +1,5 @@ import os +import contextlib from .. import documents from .. import results @@ -138,7 +139,11 @@ def _create_reader(numbering, content_types, relationships, docx_file): content_type = content_types.find_content_type(image_path) def open_image(): - return docx_file.open(image_path) + image_file = docx_file.open(image_path) + if hasattr(image_file, "__exit__"): + return image_file + else: + return contextlib.closing(image_file) image = documents.image(alt_text=alt_text, content_type=content_type, open=open_image) return results.success(image) diff --git a/scripts/mammoth b/scripts/mammoth index 16563e9..4f404ad 100644 --- a/scripts/mammoth +++ b/scripts/mammoth @@ -3,6 +3,8 @@ import argparse import contextlib import sys +import os +import shutil import mammoth @@ -16,15 +18,46 @@ def main(): with open(args.styles) as styles_fileobj: styles = styles_fileobj.read() - with _open_output(args.output) as output: - with open(args.path, "rb") as docx_fileobj: - result = mammoth.convert_to_html(docx_fileobj, styles=styles) - for message in result.messages: - sys.stderr.write(message.message) - sys.stderr.write("\n") + with open(args.path, "rb") as docx_fileobj: + if args.output_dir is None: + convert_image = None + output_path = args.output + else: + convert_image = ImageWriter(args.output_dir) + output_filename = "{0}.html".format(os.path.basename(args.path).rpartition(".")[0]) + output_path = os.path.join(args.output_dir, output_filename) + + result = mammoth.convert_to_html(docx_fileobj, styles=styles, convert_image=convert_image) + for message in result.messages: + sys.stderr.write(message.message) + sys.stderr.write("\n") + + with _open_output(output_path) as output: output.write(result.value) +class ImageWriter(object): + def __init__(self, output_dir): + self._output_dir = output_dir + self._image_number = 1 + + def __call__(self, element, html_generator): + + extension = element.content_type.partition("/")[2] + image_filename = "{0}.{1}".format(self._image_number, extension) + with open(os.path.join(self._output_dir, image_filename), "wb") as image_dest: + with element.open() as image_source: + shutil.copyfileobj(image_source, image_dest) + + attributes = {"src": image_filename} + if element.alt_text is not None: + attributes["alt"] = element.alt_text + + html_generator.self_closing("img", attributes) + + self._image_number += 1 + + @contextlib.contextmanager def _open_output(name): if name is None: @@ -52,7 +85,11 @@ class Output(object): def _parse_args(): parser = argparse.ArgumentParser() parser.add_argument("path") - parser.add_argument("output", nargs="?") + + output_group = parser.add_mutually_exclusive_group() + output_group.add_argument("output", nargs="?") + output_group.add_argument("--output-dir", nargs="?") + parser.add_argument("--styles", required=False) return parser.parse_args() diff --git a/tests/cli_tests.py b/tests/cli_tests.py index c274ce2..b23ae9a 100644 --- a/tests/cli_tests.py +++ b/tests/cli_tests.py @@ -34,4 +34,16 @@ def inline_images_are_included_in_output_if_writing_to_single_file(): docx_path = test_path("tiny-picture.docx") result = _local.run(["mammoth", docx_path]) assert_equal(b"""<p><img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAIAAAACUFjqAAAAAXNSR0IArs4c6QAAAAlwSFlzAAAOvgAADr4B6kKxwAAAABNJREFUKFNj/M+ADzDhlWUYqdIAQSwBE8U+X40AAAAASUVORK5CYII=" /></p>""", result.output) + + +@istest +def images_are_written_to_separate_files_if_output_dir_is_set(): + with tempman.create_temp_dir() as temp_dir: + output_path = os.path.join(temp_dir.path, "tiny-picture.html") + docx_path = test_path("tiny-picture.docx") + result = _local.run(["mammoth", docx_path, "--output-dir", temp_dir.path]) + assert_equal(b"", result.stderr_output) + assert_equal(b"", result.output) + with open(output_path) as output_file: + assert_equal("""<p><img src="1.png" /></p>""", output_file.read()) |