(#4) Refactor: Dataset abstraction

betarixm · Mar 23, 2022 · e5d8de2 · e5d8de2
1 parent 71fc88d
commit e5d8de2
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/src/typings/dataset.py b/src/typings/dataset.py
@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+from typing import NewType
+
+import tensorflow as tf
+
+TestSet = NewType("TestSet", tf.data.Dataset)
+TrainSet = NewType("TrainSet", tf.data.Dataset)
+
+
+class Dataset(ABC):
+    @staticmethod
+    @abstractmethod
+    def dataset() -> (TestSet, TrainSet):
+        pass
+
+    def __new__(cls, *args, **kwargs):
+        return cls.dataset()