diff --git a/webviz_config/webviz_storage/__init__.py b/webviz_config/webviz_storage/__init__.py index 023c9030..9bade47c 100644 --- a/webviz_config/webviz_storage/__init__.py +++ b/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,12 +79,10 @@ 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): @@ -93,18 +90,8 @@ def save_data(data, path: str): @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()