snack_protocol

.gitignore

@@ -157,4 +157,4 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/

src/data_snack/wrap/base.py

@@ -1,17 +1,36 @@
-from abc import ABC
-from typing import Text, List, Optional, Any
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Text, List, Optional, Any, Type, Protocol
 
 from data_snack.entities import Entity
 
 
+class SnackProtocol(Protocol):
+    def set(self, entity: Entity) -> Optional[Text]:
+        ...
+
+    def get(self, cls: Type[Entity], key_values: List[Text]) -> Entity:
+        ...
+
+    def get_many(self, cls: Type[Entity], keys_values: List[List[Text]]) -> List[Entity]:
+        ...
+
+    def set_many(self, entities: List[Entity]) -> List[Text]:
+        ...
+
+    def keys(self, cls: Type[Entity]) -> List[bytes]:
+        ...
+
+
+@dataclass
 class Wrap(ABC):
     """
     Wraps are used to provide a simplified interface for accessing `Snack` for one, selected type of `Entity`.
     """
+    snack: SnackProtocol
+    entity_type: Type[Entity]
 
-    def __init__(self, *args: Any, **kwargs: Any):
-        ...
-
+    @abstractmethod
     def set(self, entity: Entity) -> Optional[Text]:
         """
         Saves given entity in db.
@@ -21,6 +40,7 @@ def set(self, entity: Entity) -> Optional[Text]:
         """
         ...
 
+    @abstractmethod
     def get(self, key_values: List[Text]) -> Entity:
         """
         Reads entities from db based on provided key values.
@@ -30,6 +50,7 @@ def get(self, key_values: List[Text]) -> Entity:
         """
         ...
 
+    @abstractmethod
     def get_many(self, keys_values: List[List[Text]]) -> List[Entity]:
         """
         Gets list of `Entity` objects from db based on provided list of keys.
@@ -39,6 +60,7 @@ def get_many(self, keys_values: List[List[Text]]) -> List[Entity]:
         """
         ...
 
+    @abstractmethod
     def set_many(self, entities: List[Entity]) -> List[Text]:
         """
         Saves multiple `Entity` objects in db.
@@ -48,6 +70,7 @@ def set_many(self, entities: List[Entity]) -> List[Text]:
         """
         ...
 
+    @abstractmethod
     def keys(self) -> List[bytes]:
         """
         Reads a list of keys available in db for given `Entity` type.

src/data_snack/wrap/entity.py

@@ -7,8 +7,6 @@
 
 @dataclass
 class EntityWrap(Wrap):
-    snack: "Snack"
-    entity_type: Type[Entity]
     _entity_type_name: Text = field(init=False)
 
     @property

tests/data_snack/conftest.py

@@ -2,7 +2,7 @@
 from unittest.mock import MagicMock
 
 import pytest
-from pydantic.dataclasses import dataclass
+from dataclasses import dataclass
 
 from data_snack import Snack
 from data_snack.connections import Connection