From e47467856cfed0a4014c68575bfb0dd11c1a8f7f Mon Sep 17 00:00:00 2001
From: Hans Kallekleiv <hakal@statoil.com>
Date: Sun, 31 Jan 2021 11:05:49 +0100
Subject: [PATCH] wip

---
 webviz_config/webviz_storage/__init__.py | 40 +++++++++++++-----------
 1 file changed, 22 insertions(+), 18 deletions(-)

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())