Allow io.BytesIO as webviz storage type (#180)

CONTRIBUTING.md

@@ -367,7 +367,7 @@ The core of `webviz-config` will do the following:
       functions will return the dataframe from the stored `.parquet` files,
       instead of running the actual function code.
 
-Currently only functions returning `pd.DataFrame` and `pathlib.Path` are
+Currently only functions returning `pd.DataFrame`, `io.BytesIO` or `pathlib.Path` are
 supported. A `NotImplementedError` will be raised if the function being
 decorated with `@webvizstore` does not have a type hint being supported.
 

webviz_config/webviz_store.py

@@ -1,4 +1,5 @@
 import os
+import io
 import glob
 import shutil
 import functools
@@ -13,6 +14,9 @@
 
 
 class WebvizStorage:
+
+    RETURN_TYPES = [pd.DataFrame, pathlib.Path, io.BytesIO]
+
     def __init__(self):
         self._use_storage = False
         self.storage_functions = set()
@@ -25,10 +29,9 @@ def register_function(self, func):
 
         return_type = inspect.getfullargspec(func).annotations["return"]
 
-        if return_type not in [pd.DataFrame, pathlib.Path]:
+        if return_type not in WebvizStorage.RETURN_TYPES:
             raise NotImplementedError(
-                "Currently only storage of dataframes "
-                "and file resources are implemented."
+                f"Webviz storage type must be one of {WebvizStorage.RETURN_TYPES}"
             )
 
         self.storage_functions.add(func)
@@ -140,6 +143,8 @@ def get_stored_data(self, func, *args, **kwargs):
                 return pd.read_parquet(f"{path}.parquet")
             if return_type == pathlib.Path:
                 return pathlib.Path(glob.glob(f"{path}*")[0])
+            if return_type == io.BytesIO:
+                return io.BytesIO(pathlib.Path(path).read_bytes())
             raise ValueError(f"Unknown return type {return_type}")
 
         except OSError:
@@ -175,6 +180,8 @@ def build_store(self):
                     output.to_parquet(f"{path}.parquet")
                 elif isinstance(output, pathlib.Path):
                     shutil.copy(output, f"{path}{output.suffix}")
+                elif isinstance(output, io.BytesIO):
+                    pathlib.Path(path).write_bytes(output.getvalue())
                 else:
                     raise ValueError(f"Unknown return type {type(output)}")