wip

webviz_config/webviz_storage/__init__.py

@@ -30,13 +30,12 @@ class WebvizStorageType(abc.ABC):
     @staticmethod
     @abc.abstractmethod
     def get_data(path: str, **kwargs):
-        """ Abstract method to run a command """
-        pass
+        """ Abstract method to retrieve stored data """
 
     @staticmethod
     @abc.abstractmethod
     def save_data(data, path: str):
-        pass
+        """ Abstract method to save data to store """
 
 
 class WebvizStorageTypeRegistry:
@@ -55,7 +54,7 @@ def inner_wrapper(wrapped_class: WebvizStorageType) -> Callable:
         return inner_wrapper
 
     @classmethod
-    def create_storage_type(cls, return_type, **kwargs) -> "WebvizStorageType":
+    def create_storage_type(cls, return_type: str, **kwargs) -> "WebvizStorageType":
         """Factory command to create the storage type.
         This method gets the appropriate WebvizStorageType class from the registry
         and creates an instance of it, while passing in the parameters
@@ -80,31 +79,19 @@ def return_types(cls):
 
 
 @WebvizStorageTypeRegistry.register(pd.DataFrame)
-class DataFrame(WebvizStorageType):
+class TypeDataFrame(WebvizStorageType):
     @staticmethod
-    def get_data(path, load_args):
-        return pd.read_parquet(
-            f"{path}.parquet", columns=load_args.get("parquet_columns", None)
-        )
+    def get_data(path, **kwargs):
+        return pd.read_parquet(f"{path}.parquet")
 
     @staticmethod
     def save_data(data, path: str):
         data.to_parquet(f"{path}.parquet")
 
 
 @WebvizStorageTypeRegistry.register(pathlib.Path)
-class Path(WebvizStorageType):
-    @staticmethod
-    def get_data(path, **kwargs):
-        return pathlib.Path(glob.glob(f"{path}*")[0])
-
-    @staticmethod
-    def save_data(data, path: str):
-        shutil.copy(data, f"{path}{data.suffix}")
-
-
 @WebvizStorageTypeRegistry.register(pathlib.PosixPath)
-class PosixPath(WebvizStorageType):
+class TypePath(WebvizStorageType):
     @staticmethod
     def get_data(path, **kwargs):
         return pathlib.Path(glob.glob(f"{path}*")[0])
@@ -115,18 +102,8 @@ def save_data(data, path: str):
 
 
 @WebvizStorageTypeRegistry.register(list)
-class TypingList(WebvizStorageType):
-    @staticmethod
-    def get_data(path, **kwargs):
-        return np.load(f"{path}.npy").tolist()
-
-    @staticmethod
-    def save_data(data, path: str):
-        np.save(f"{path}.npy", data)
-
-
 @WebvizStorageTypeRegistry.register(List)
-class TypingList(WebvizStorageType):
+class TypeList(WebvizStorageType):
     @staticmethod
     def get_data(path, **kwargs):
         return np.load(f"{path}.npy").tolist()
@@ -137,7 +114,7 @@ def save_data(data, path: str):
 
 
 @WebvizStorageTypeRegistry.register(io.BytesIO)
-class BytesIO(WebvizStorageType):
+class TypeBytesIO(WebvizStorageType):
     @staticmethod
     def get_data(path, **kwargs):
         return np.load(f"{path}.npy").tolist()