Add json_serializable keyword

compute-tooling · Sep 25, 2019 · 499f58a · 499f58a
1 parent 979474d
commit 499f58a
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 12 deletions.
diff --git a/cs_storage/__init__.py b/cs_storage/__init__.py
@@ -24,29 +24,32 @@ def __init__(self, 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):
-        return base64.b64encode(data).decode("utf-8")
+    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"))
@@ -139,7 +142,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()
@@ -154,7 +157,7 @@ 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(
                 {
                     "title": rem_output["title"],

diff --git a/cs_storage/tests/test_cs_storage.py b/cs_storage/tests/test_cs_storage.py
@@ -106,7 +106,7 @@ def test_get_serializer():
         assert cs_storage.get_serializer(t)
 
 
-def test_cs_storage():
+def test_cs_storage(png, jpg):
     exp_loc_res = {
         "renderable": [
             {
@@ -122,12 +122,12 @@ def test_cs_storage():
             {
                 "media_type": "PNG",
                 "title": "PNG data",
-                "data": b"PNG bytes",
+                "data": png,
             },
             {
                 "media_type": "JPEG",
                 "title": "JPEG data",
-                "data": b"JPEG bytes",
+                "data": jpg,
             },
             {
                 "media_type": "MP3",
@@ -171,11 +171,17 @@ def test_cs_storage():
     }
     task_id = uuid.uuid4()
     rem_res = cs_storage.write(task_id, exp_loc_res)
-    loc_res = cs_storage.read(rem_res)
+    loc_res = cs_storage.read(rem_res, json_serializable=False)
     assert loc_res == exp_loc_res
+    assert json.dumps(
+        cs_storage.read(rem_res, json_serializable=True)
+    )
 
-    loc_res1 = cs_storage.read({"renderable": rem_res["renderable"]})
+    loc_res1 = cs_storage.read({"renderable": rem_res["renderable"]}, json_serializable=False)
     assert loc_res1["renderable"] == exp_loc_res["renderable"]
+    assert json.dumps(
+        cs_storage.read({"renderable": rem_res["renderable"]}, json_serializable=True)
+    )
 
 
 def test_errors():