Fix handling of PNG/JPG data

cs_storage/__init__.py

@@ -1,3 +1,4 @@
+import base64
 import io
 import json
 import os
@@ -43,17 +44,25 @@ def deserialize(self, data):
         return data.decode()
 
 
+class Base64Serializer(Serializer):
+    def deserialize(self, data):
+        return base64.b64encode(data).decode("utf-8")
+
+    def from_string(self, data):
+        return base64.b64decode(data.encode("utf-8"))
+
+
 def get_serializer(media_type):
     return {
         "bokeh": JSONSerializer("json"),
         "table": TextSerializer("html"),
         "CSV": TextSerializer("csv"),
-        "PNG": Serializer("png"),
-        "JPEG": Serializer("jpeg"),
-        "MP3": Serializer("mp3"),
-        "MP4": Serializer("mp4"),
-        "HDF5": Serializer("h5"),
-        "PDF": Serializer("pdf"),
+        "PNG": Base64Serializer("png"),
+        "JPEG": Base64Serializer("jpeg"),
+        "MP3": Base64Serializer("mp3"),
+        "MP4": Base64Serializer("mp4"),
+        "HDF5": Base64Serializer("h5"),
+        "PDF": Base64Serializer("pdf"),
         "Markdown": TextSerializer("md"),
         "Text": TextSerializer("txt"),
     }[media_type]

cs_storage/tests/test_cs_storage.py

@@ -1,3 +1,4 @@
+import base64
 import io
 import json
 import uuid
@@ -10,6 +11,44 @@
 import cs_storage
 
 
+@pytest.fixture
+def png():
+    import matplotlib.pyplot as plt
+    import numpy as np
+    x = np.linspace(0, 2, 100)
+    plt.figure()
+    plt.plot(x, x, label='linear')
+    plt.plot(x, x**2, label='quadratic')
+    plt.plot(x, x**3, label='cubic')
+    plt.xlabel('x label')
+    plt.ylabel('y label')
+    plt.title("Simple Plot")
+    plt.legend()
+    initial_buff = io.BytesIO()
+    plt.savefig(initial_buff, format="png")
+    initial_buff.seek(0)
+    return initial_buff.read()
+
+
+@pytest.fixture
+def jpg():
+    import matplotlib.pyplot as plt
+    import numpy as np
+    x = np.linspace(0, 2, 100)
+    plt.figure()
+    plt.plot(x, x, label='linear')
+    plt.plot(x, x**2, label='quadratic')
+    plt.plot(x, x**3, label='cubic')
+    plt.xlabel('x label')
+    plt.ylabel('y label')
+    plt.title("Simple Plot")
+    plt.legend()
+    initial_buff = io.BytesIO()
+    plt.savefig(initial_buff, format="jpg")
+    initial_buff.seek(0)
+    return initial_buff.read()
+
+
 def test_JSONSerializer():
     ser = cs_storage.JSONSerializer("json")
 
@@ -46,6 +85,21 @@ def test_serializer():
     assert act == b"hello world"
 
 
+def test_base64serializer(png, jpg):
+    """Test round trip serializtion/deserialization of PNG and JPG"""
+    ser = cs_storage.Base64Serializer("PNG")
+    asbytes = ser.serialize(png)
+    asstr = ser.deserialize(asbytes)
+    assert png == ser.from_string(asstr)
+    assert json.dumps({"pic": asstr})
+
+    ser = cs_storage.Base64Serializer("JPG")
+    asbytes = ser.serialize(jpg)
+    asstr = ser.deserialize(asbytes)
+    assert jpg == ser.from_string(asstr)
+    assert json.dumps({"pic": asstr})
+
+
 def test_get_serializer():
     types = ["bokeh", "table", "CSV", "PNG", "JPEG", "MP3", "MP4", "HDF5"]
     for t in types:

environment.yml

@@ -5,3 +5,5 @@ dependencies:
   - "marshmallow>=3.0.0"
   - pytest
   - gcsfs
+  - matplotlib
+  - numpy