shared/datastore/src/datastore_interface.py

@@ -0,0 +1,66 @@
+"""Definition of datastore interface."""
+
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+
+from PIL import Image
+from torch.nn import Module
+
+
+class DatastoreInterface(ABC):
+    """Datastore interface class."""
+
+    @abstractmethod
+    def connect(
+        self,
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def close(
+        self,
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def __enter__(
+        self,
+    ) -> None:
+        self.connect()
+
+    @abstractmethod
+    def __exit__(
+        self,
+        exc_type,
+        exc_val,
+        exc_tb,
+    ) -> None:
+        self.close()
+
+    @abstractmethod
+    def put_image(
+        self,
+        image: Image.Image,
+    ) -> str:
+        pass
+
+    @abstractmethod
+    def get_image(
+        self,
+        object_name: str,
+    ) -> Image.Image:
+        pass
+
+    @abstractmethod
+    def put_model(
+        self,
+        model: Module,
+    ) -> str:
+        pass
+
+    @abstractmethod
+    def get_model(
+        self,
+        object_name: str,
+    ) -> OrderedDict:
+        pass