Extract common image conversion code

4 files changed, 20 insertions, 18 deletions

diff --git a/mammoth/__init__.py b/mammoth/__init__.py
index 05e2167..0b524b1 100644
--- a/mammoth/__init__.py
+++ b/mammoth/__init__.py
@@ -1,5 +1,7 @@
 from .results import Result
-from . import docx, conversion, options, documents
+from . import docx, conversion, options, documents, images
+
+__all__ = ["convert_to_html", "extract_raw_text", "images"]
 
 
 def convert_to_html(fileobj, transform_document=None, **kwargs):
diff --git a/mammoth/conversion.py b/mammoth/conversion.py
index 0349774..715e5e0 100644
--- a/mammoth/conversion.py
+++ b/mammoth/conversion.py
@@ -1,7 +1,7 @@
 import base64
 import io
 
-from . import documents, results, html_paths
+from . import documents, results, html_paths, images
 from .html_generation import HtmlGenerator, satisfy_html_path
 
 
@@ -29,7 +29,7 @@ class DocumentConverter(object):
             documents.Table: self._convert_table,
             documents.TableRow: self._convert_table_row,
             documents.TableCell: self._convert_table_cell,
-            documents.Image: convert_image or self._convert_image,
+            documents.Image: convert_image or images.inline(self._convert_image),
         }
 
 
@@ -99,17 +99,13 @@ class DocumentConverter(object):
         html_generator.end()
     
     
-    def _convert_image(self, image, html_generator):
+    def _convert_image(self, image):
         with image.open() as image_bytes:
             encoded_src = base64.b64encode(image_bytes.read()).decode("ascii")
         
-        attributes = {
+        return {
             "src": "data:{0};base64,{1}".format(image.content_type, encoded_src)
         }
-        if image.alt_text:
-            attributes["alt"] = image.alt_text
-            
-        html_generator.self_closing("img", attributes)
 
 
     def _convert_elements_to_html(self, elements, html_generator):
diff --git a/mammoth/images.py b/mammoth/images.py
new file mode 100644
index 0000000..713f739
--- /dev/null
+++ b/mammoth/images.py
@@ -0,0 +1,9 @@
+def inline(func):
+    def convert_image(image, html_generator):
+        attributes = func(image).copy()
+        if image.alt_text:
+            attributes["alt"] = image.alt_text
+            
+        html_generator.self_closing("img", attributes)
+    
+    return convert_image
diff --git a/scripts/mammoth b/scripts/mammoth
index 1b2fed1..b6fa056 100644
--- a/scripts/mammoth
+++ b/scripts/mammoth
@@ -23,7 +23,7 @@ def main():
             convert_image = None
             output_path = args.output
         else:
-            convert_image = ImageWriter(args.output_dir)
+            convert_image = mammoth.images.inline(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)
         
@@ -45,21 +45,16 @@ class ImageWriter(object):
         self._output_dir = output_dir
         self._image_number = 1
         
-    def __call__(self, element, html_generator):
-        
+    def __call__(self, element):
         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
+        
+        return {"src": image_filename}
 
 
 @contextlib.contextmanager