From 2c9f6406977ea04d02d1ed6e54035ec2c593417d Mon Sep 17 00:00:00 2001
From: Maciej Odziemczyk <maciej.odziemczyk@webinterpret.com>
Date: Wed, 21 Dec 2022 14:36:57 +0100
Subject: [PATCH] add: snack protocol

---
 .gitignore                    |  2 +-
 src/data_snack/wrap/base.py   | 33 ++++++++++++++++++++++++++++-----
 src/data_snack/wrap/entity.py |  2 --
 tests/data_snack/conftest.py  |  2 +-
 4 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/.gitignore b/.gitignore
index df3768d..e020613 100644
--- a/.gitignore
+++ b/.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/
diff --git a/src/data_snack/wrap/base.py b/src/data_snack/wrap/base.py
index 6a29aa3..2af1c40 100644
--- a/src/data_snack/wrap/base.py
+++ b/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.
diff --git a/src/data_snack/wrap/entity.py b/src/data_snack/wrap/entity.py
index f6198c4..7a34287 100644
--- a/src/data_snack/wrap/entity.py
+++ b/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
diff --git a/tests/data_snack/conftest.py b/tests/data_snack/conftest.py
index df5302c..ff7c483 100644
--- a/tests/data_snack/conftest.py
+++ b/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