Require pillow<7 until new matplotlib comes out (>3.2.1)

cs_storage/__init__.py

@@ -1,3 +1,4 @@
+import base64
 import io
 import json
 import os
@@ -12,50 +13,65 @@
 
 from .screenshot import screenshot, ScreenshotError, SCREENSHOT_ENABLED
 
-__version__ = "1.6.0"
+__version__ = "1.7.0"
 
 
 BUCKET = os.environ.get("BUCKET", None)
 
 
 class Serializer:
+    """
+    Base class for serializng input data to bytes and back.
+    """
+
     def __init__(self, ext):
         self.ext = ext
 
     def serialize(self, data):
         return data
 
-    def deserialize(self, data):
+    def deserialize(self, data, json_serializable=True):
         return data
 
 
 class JSONSerializer(Serializer):
     def serialize(self, data):
         return json.dumps(data).encode()
 
-    def deserialize(self, data):
+    def deserialize(self, data, json_serializable=True):
         return json.loads(data.decode())
 
 
 class TextSerializer(Serializer):
     def serialize(self, data):
         return data.encode()
 
-    def deserialize(self, data):
+    def deserialize(self, data, json_serializable=True):
         return data.decode()
 
 
+class Base64Serializer(Serializer):
+    def deserialize(self, data, json_serializable=True):
+        if json_serializable:
+            return base64.b64encode(data).decode("utf-8")
+        else:
+            return data
+
+    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]
@@ -173,7 +189,7 @@ def write(task_id, loc_result, do_upload=True):
     return rem_result
 
 
-def read(rem_result):
+def read(rem_result, json_serializable=True):
     # compute studio results have public read access.
     fs = gcsfs.GCSFileSystem(token="anon")
     s = time.time()
@@ -188,7 +204,9 @@ def read(rem_result):
 
         for rem_output in rem_result[category]["outputs"]:
             ser = get_serializer(rem_output["media_type"])
-            rem_data = ser.deserialize(zipfileobj.read(rem_output["filename"]))
+            rem_data = ser.deserialize(
+                zipfileobj.read(rem_output["filename"]), json_serializable
+            )
             read[category].append(
                 {
                     "id": rem_output.get("id", None),

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,13 +85,28 @@ 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:
         assert cs_storage.get_serializer(t)
 
 
-def test_cs_storage():
+def test_cs_storage(png, jpg):
     dummy_uuid = "c7a65ad2-0c2c-45d7-b0f7-d9fd524c49b3"
     exp_loc_res = {
         "renderable": [
@@ -61,11 +115,32 @@ def test_cs_storage():
                 "title": "bokeh plot",
                 "data": {"html": "<div/>", "javascript": "console.log('hello world')"},
             },
-            {"media_type": "table", "title": "table stuff", "data": "<table/>"},
-            {"media_type": "PNG", "title": "PNG data", "data": b"PNG bytes"},
-            {"media_type": "JPEG", "title": "JPEG data", "data": b"JPEG bytes"},
-            {"media_type": "MP3", "title": "MP3 data", "data": b"MP3 bytes"},
-            {"media_type": "MP4", "title": "MP4 data", "data": b"MP4 bytes"},
+            {
+                "media_type": "table",
+                "title": "table stuff",
+                "data": "<table/>",
+            },
+            {
+                "media_type": "PNG",
+                "title": "PNG data",
+                "data": png,
+            },
+            {
+                "media_type": "JPEG",
+                "title": "JPEG data",
+                "data": jpg,
+            },
+            {
+                "media_type": "MP3",
+                "title": "MP3 data",
+                "data": b"MP3 bytes",
+            },
+
+            {
+                "media_type": "MP4",
+                "title": "MP4 data",
+                "data": b"MP4 bytes",
+            },
         ],
         "downloadable": [
             {"media_type": "CSV", "title": "CSV file", "data": "comma,sep,values\n"},

environment.yml

@@ -9,3 +9,6 @@ dependencies:
   - pyppeteer # optional
   - bokeh # optional
   - websockets # optional
+  - matplotlib
+  - numpy
+  - "pillow<7"