Add --output-dir option

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())