diff --git a/webviz_config/webviz_storage/__init__.py b/webviz_config/webviz_storage/__init__.py index 9bade47c..4fd681e8 100644 --- a/webviz_config/webviz_storage/__init__.py +++ b/webviz_config/webviz_storage/__init__.py @@ -8,7 +8,7 @@ import pathlib import warnings from collections import defaultdict -from typing import Callable, List, Union, Any +from typing import Callable, List, Union, Any, Type, Dict import weakref import numpy as np @@ -17,10 +17,12 @@ class ClassProperty: - def __init__(self, fget): + def __init__(self, fget: Callable): self.fget = fget - def __get__(self, owner_self, owner_cls): + def __get__(self, owner_self: Type[Any], owner_cls: Type[Any]) -> Any: + print(type(owner_self), type(owner_cls)) + print(type(self.fget(owner_cls))) return self.fget(owner_cls) @@ -29,23 +31,23 @@ class WebvizStorageType(abc.ABC): @staticmethod @abc.abstractmethod - def get_data(path: str, **kwargs): + def get_data(path: str, **kwargs: Dict) -> None: """ Abstract method to retrieve stored data """ @staticmethod @abc.abstractmethod - def save_data(data, path: str): + def save_data(data: Any, path: str) -> Any: """ Abstract method to save data to store """ class WebvizStorageTypeRegistry: """ Registry of allowed webviz storage types """ - registry = {} + registry: Dict = {} @classmethod - def register(cls, return_type): - def inner_wrapper(wrapped_class: WebvizStorageType) -> Callable: + def register(cls, return_type: Type) -> Callable: + def inner_wrapper(wrapped_class: Type) -> Type: if return_type in cls.registry: print(f"Storage type {return_type} already exists. Will replace it") cls.registry[return_type] = wrapped_class @@ -54,7 +56,9 @@ def inner_wrapper(wrapped_class: WebvizStorageType) -> Callable: return inner_wrapper @classmethod - def create_storage_type(cls, return_type: str, **kwargs) -> "WebvizStorageType": + def create_storage_type( + cls, return_type: str, **kwargs: Dict + ) -> Union[None, 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 @@ -74,18 +78,18 @@ def create_storage_type(cls, return_type: str, **kwargs) -> "WebvizStorageType": # pylint: disable=no-self-argument @ClassProperty - def return_types(cls): + def return_types(cls) -> List: return list(cls.registry.keys()) @WebvizStorageTypeRegistry.register(pd.DataFrame) class TypeDataFrame(WebvizStorageType): @staticmethod - def get_data(path, **kwargs): + def get_data(path: str, **kwargs: Dict) -> Any: return pd.read_parquet(f"{path}.parquet") @staticmethod - def save_data(data, path: str): + def save_data(data: Any, path: str) -> None: data.to_parquet(f"{path}.parquet") @@ -93,11 +97,11 @@ def save_data(data, path: str): @WebvizStorageTypeRegistry.register(pathlib.PosixPath) class TypePath(WebvizStorageType): @staticmethod - def get_data(path, **kwargs): + def get_data(path: str, **kwargs: Dict) -> Any: return pathlib.Path(glob.glob(f"{path}*")[0]) @staticmethod - def save_data(data, path: str): + def save_data(data: Any, path: str) -> None: shutil.copy(data, f"{path}{data.suffix}") @@ -105,20 +109,20 @@ def save_data(data, path: str): @WebvizStorageTypeRegistry.register(List) class TypeList(WebvizStorageType): @staticmethod - def get_data(path, **kwargs): + def get_data(path: str, **kwargs: Dict) -> Any: return np.load(f"{path}.npy").tolist() @staticmethod - def save_data(data, path: str): + def save_data(data: Any, path: str) -> None: np.save(f"{path}.npy", data) @WebvizStorageTypeRegistry.register(io.BytesIO) class TypeBytesIO(WebvizStorageType): @staticmethod - def get_data(path, **kwargs): + def get_data(path: str, **kwargs: Dict) -> Any: return np.load(f"{path}.npy").tolist() @staticmethod - def save_data(data, path: str): + def save_data(data: Any, path: str) -> None: pathlib.Path(path).write_bytes(data.getvalue())