Handle image encoding properly

4 files changed, 13 insertions, 8 deletions

diff --git a/mammoth/conversion.py b/mammoth/conversion.py
index 87e1bf9..b3b41d0 100644
--- a/mammoth/conversion.py
+++ b/mammoth/conversion.py
@@ -74,9 +74,7 @@ class DocumentConverter(object):
     
     def _convert_image(self, image, html_generator):
         with image.open() as image_bytes:
-            output = io.BytesIO()
-            base64.encode(image_bytes, output)
-            encoded_src = output.getvalue().rstrip("\n")
+            encoded_src = base64.b64encode(image_bytes.read()).decode("ascii")
         
         attributes = {
             "src": "data:{0};base64,{1}".format(image.content_type, encoded_src)
diff --git a/mammoth/html_generation.py b/mammoth/html_generation.py
index 563e2d6..9716ac2 100644
--- a/mammoth/html_generation.py
+++ b/mammoth/html_generation.py
@@ -1,4 +1,5 @@
 import cgi
+import sys
 
 
 class HtmlGenerator(object):
@@ -67,7 +68,7 @@ def _generate_attribute_string(attributes):
     else:
         return "".join(
             ' {0}="{1}"'.format(key, _escape_html(value))
-            for key, value in attributes.iteritems()
+            for key, value in _iteritems(attributes)
         )
 
 
@@ -101,3 +102,9 @@ def _is_element_match(generated_element, path_element):
 
 def _generate_class_attribute(path_element):
     return " ".join(path_element.class_names)
+
+
+if sys.version_info[0] == 2:
+    _iteritems = lambda x: x.iteritems()
+else:
+    _iteritems = lambda x: x.items()
diff --git a/tests/conversion_tests.py b/tests/conversion_tests.py
index d9742cc..cb9b064 100644
--- a/tests/conversion_tests.py
+++ b/tests/conversion_tests.py
@@ -145,14 +145,14 @@ def docx_tab_is_converted_to_tab_in_html():
 
 @istest
 def images_are_converted_to_img_tags_with_data_uri():
-    image = documents.image(alt_text=None, content_type="image/png", open=lambda: io.BytesIO("abc"))
+    image = documents.image(alt_text=None, content_type="image/png", open=lambda: io.BytesIO(b"abc"))
     result = convert_document_element_to_html(image)
     assert_equal('<img src="data:image/png;base64,YWJj" />', result.value)
 
 
 @istest
 def images_have_alt_tags_if_available():
-    image = documents.image(alt_text="It's a hat", content_type="image/png", open=lambda: io.BytesIO("abc"))
+    image = documents.image(alt_text="It's a hat", content_type="image/png", open=lambda: io.BytesIO(b"abc"))
     result = convert_document_element_to_html(image)
     assert_equal('<img src="data:image/png;base64,YWJj" alt="It\'s a hat" />', result.value)
 
@@ -163,7 +163,7 @@ def can_define_custom_conversion_for_images():
         with image.open() as image_file:
             html_generator.self_closing("img", {"alt": image_file.read()})
         
-    image = documents.image(alt_text=None, content_type="image/png", open=lambda: io.BytesIO("abc"))
+    image = documents.image(alt_text=None, content_type="image/png", open=lambda: io.BytesIO(b"abc"))
     result = convert_document_element_to_html(image, convert_image=convert_image)
     assert_equal('<img alt="abc" />', result.value)
 
diff --git a/tests/mammoth_tests.py b/tests/mammoth_tests.py
index bfa88b5..f2469a9 100644
--- a/tests/mammoth_tests.py
+++ b/tests/mammoth_tests.py
@@ -17,7 +17,7 @@ def docx_containing_one_paragraph_is_converted_to_single_p_element():
 def inline_images_are_included_in_output():
     with open(test_path("tiny-picture.docx"), "rb") as fileobj:
         result = mammoth.convert_to_html(fileobj=fileobj)
-        assert_equal("""<p><img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAIAAAACUFjqAAAAAXNSR0IArs4c6QAAAAlwSFlzAAAO\nvgAADr4B6kKxwAAAABNJREFUKFNj/M+ADzDhlWUYqdIAQSwBE8U+X40AAAAASUVORK5CYII=" /></p>""", result.value)
+        assert_equal("""<p><img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAIAAAACUFjqAAAAAXNSR0IArs4c6QAAAAlwSFlzAAAOvgAADr4B6kKxwAAAABNJREFUKFNj/M+ADzDhlWUYqdIAQSwBE8U+X40AAAAASUVORK5CYII=" /></p>""", result.value)
         assert_equal([], result.messages)