Files
sillyhome-future/app/core/stores.py

75 lines
2.3 KiB
Python

from __future__ import annotations
import json
import os
from pathlib import Path
from threading import RLock
from typing import TypeVar
from pydantic import BaseModel
from app.core.models import BackupBundle, ControlState, LearningState, RuntimeState
T = TypeVar("T", bound=BaseModel)
class JsonDocumentStore:
def __init__(self, root: str | Path) -> None:
self.root = Path(root).resolve()
self.root.mkdir(parents=True, exist_ok=True)
self._lock = RLock()
def load(self, name: str, model: type[T], default: T) -> T:
path = self.root / name
with self._lock:
if not path.exists():
return default
return model.model_validate_json(path.read_text(encoding="utf-8"))
def save(self, name: str, value: T) -> T:
path = self.root / name
temporary = path.with_suffix(path.suffix + ".tmp")
with self._lock:
temporary.write_text(
json.dumps(value.model_dump(mode="json"), ensure_ascii=True, sort_keys=True) + "\n",
encoding="utf-8",
)
os.replace(temporary, path)
return value
class FutureStores:
def __init__(self, root: str | Path = ".future_store") -> None:
self._docs = JsonDocumentStore(root)
def runtime(self) -> RuntimeState:
return self._docs.load("runtime.json", RuntimeState, RuntimeState())
def save_runtime(self, state: RuntimeState) -> RuntimeState:
return self._docs.save("runtime.json", state)
def learning(self) -> LearningState:
return self._docs.load("learning.json", LearningState, LearningState())
def save_learning(self, state: LearningState) -> LearningState:
return self._docs.save("learning.json", state)
def control(self) -> ControlState:
return self._docs.load("control.json", ControlState, ControlState())
def save_control(self, state: ControlState) -> ControlState:
return self._docs.save("control.json", state)
def export_backup(self) -> BackupBundle:
return BackupBundle(
runtime=self.runtime(),
learning=self.learning(),
control=self.control(),
)
def restore_backup(self, bundle: BackupBundle) -> None:
self.save_runtime(bundle.runtime)
self.save_learning(bundle.learning)
self.save_control(bundle.control)