49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
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
|