Convert PIL image objects to data uri strings in JSON serialization. (plotly#1991)

jonmmease · web-flow · commit 97e4378b8c9e · 2019-12-10T08:28:36.000-05:00
This conversion is already done by the layout.image.source validator,
but this way it will also happen when serializing from a dict without
validation, and for images that show up elsewhere in the figure
(as mapbox layers for example)
diff --git a/packages/python/plotly/_plotly_utils/basevalidators.py b/packages/python/plotly/_plotly_utils/basevalidators.py
@@ -2347,20 +2347,25 @@ def validate_coerce(self, v):
             pass
         elif self._PIL and isinstance(v, self._PIL.Image.Image):
             # Convert PIL image to png data uri string
-            in_mem_file = io.BytesIO()
-            v.save(in_mem_file, format="PNG")
-            in_mem_file.seek(0)
-            img_bytes = in_mem_file.read()
-            base64_encoded_result_bytes = base64.b64encode(img_bytes)
-            base64_encoded_result_str = base64_encoded_result_bytes.decode("ascii")
-            v = "data:image/png;base64,{base64_encoded_result_str}".format(
-                base64_encoded_result_str=base64_encoded_result_str
-            )
+            v = self.pil_image_to_uri(v)
         else:
             self.raise_invalid_val(v)
 
         return v
 
+    @staticmethod
+    def pil_image_to_uri(v):
+        in_mem_file = io.BytesIO()
+        v.save(in_mem_file, format="PNG")
+        in_mem_file.seek(0)
+        img_bytes = in_mem_file.read()
+        base64_encoded_result_bytes = base64.b64encode(img_bytes)
+        base64_encoded_result_str = base64_encoded_result_bytes.decode("ascii")
+        v = "data:image/png;base64,{base64_encoded_result_str}".format(
+            base64_encoded_result_str=base64_encoded_result_str
+        )
+        return v
+
 
 class CompoundValidator(BaseValidator):
     def __init__(self, plotly_name, parent_name, data_class_str, data_docs, **kwargs):
diff --git a/packages/python/plotly/_plotly_utils/utils.py b/packages/python/plotly/_plotly_utils/utils.py
@@ -4,6 +4,7 @@
 import re
 
 from _plotly_utils.optional_imports import get_module
+from _plotly_utils.basevalidators import ImageUriValidator
 
 
 PY36_OR_LATER = sys.version_info.major == 3 and sys.version_info.minor >= 6
@@ -104,6 +105,7 @@ def default(self, obj):
             self.encode_as_date,
             self.encode_as_list,  # because some values have `tolist` do last.
             self.encode_as_decimal,
+            self.encode_as_pil,
         )
         for encoding_method in encoding_methods:
             try:
@@ -192,6 +194,15 @@ def encode_as_decimal(obj):
         else:
             raise NotEncodable
 
+    @staticmethod
+    def encode_as_pil(obj):
+        """Attempt to convert PIL.Image.Image to base64 data uri"""
+        pil = get_module("PIL")
+        if isinstance(obj, pil.Image.Image):
+            return ImageUriValidator.pil_image_to_uri(obj)
+        else:
+            raise NotEncodable
+
 
 class NotEncodable(Exception):
     pass
diff --git a/packages/python/plotly/plotly/tests/test_optional/test_utils/test_utils.py b/packages/python/plotly/plotly/tests/test_optional/test_utils/test_utils.py
@@ -16,9 +16,12 @@
 from nose.plugins.attrib import attr
 from pandas.util.testing import assert_series_equal
 import json as _json
+import os
+import base64
 
 from plotly import optional_imports, utils
 from plotly.graph_objs import Scatter, Scatter3d, Figure, Data
+from PIL import Image
 
 
 matplotlylib = optional_imports.get_module("plotly.matplotlylib")
@@ -274,6 +277,21 @@ def test_datetime_dot_date(self):
         j1 = _json.dumps(a, cls=utils.PlotlyJSONEncoder)
         assert j1 == '["2014-01-01", "2014-01-02"]'
 
+    def test_pil_image_encoding(self):
+        import _plotly_utils
+
+        img_path = os.path.join(
+            _plotly_utils.__path__[0], "tests", "resources", "1x1-black.png"
+        )
+
+        with open(img_path, "rb") as f:
+            hex_bytes = base64.b64encode(f.read()).decode("ascii")
+            expected_uri = "data:image/png;base64," + hex_bytes
+
+        img = Image.open(img_path)
+        j1 = _json.dumps({"source": img}, cls=utils.PlotlyJSONEncoder)
+        assert j1 == '{"source": "%s"}' % expected_uri
+
 
 if matplotlylib:
 
