Move pickle functions to PickleRepo

prophet/infra/improvement_pickle_repo.py

@@ -0,0 +1,49 @@
+import pickle
+from pathlib import Path
+from typing import override
+
+from prophet.domain.improvement import Improvement
+from prophet.domain.improvement_repo import IImprovementRepo, ImprovementNotFoundError
+
+
+class ImprovementPickleRepo(IImprovementRepo):
+    pickle_dir: Path
+
+    def __init__(self, pickle_dir: str | Path = "/tmp/pollenprophet") -> None:
+        self.pickle_dir = Path(pickle_dir)
+        self.pickle_dir.mkdir(parents=True, exist_ok=True)
+
+    @override
+    def add(self, improvement: Improvement) -> None:
+        fname = self.pickle_dir / improvement.id
+        try:
+            with open(fname, "wb") as f:
+                pickle.dump(improvement, f)
+                print(f"Saved {fname}")
+        except FileExistsError:
+            print(f"Error saving file {fname}")
+
+    @override
+    def add_all(self, improvements: list[Improvement]) -> None:
+        for imp in improvements:
+            self.add(imp)
+
+    @override
+    def get(self, id: str) -> Improvement:
+        try:
+            with open(self.pickle_dir / id, "rb") as f:
+                improvement: Improvement = pickle.load(f)
+        except FileNotFoundError:
+            raise ImprovementNotFoundError
+
+        return improvement
+
+    @override
+    def get_all(self) -> list[Improvement]:
+        improvements: list[Improvement] = []
+        for fname in Path(self.pickle_dir).iterdir():
+            try:
+                improvements.append(self.get(fname.name))
+            except ImprovementNotFoundError:
+                print(f"File {fname.absolute()} is not a valid Improvement.")
+        return improvements