82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from threading import RLock
|
|
from typing import cast, 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()
|
|
self._cache: dict[str, BaseModel] = {}
|
|
|
|
def load(self, name: str, model: type[T], default: T) -> T:
|
|
path = self.root / name
|
|
with self._lock:
|
|
cached = self._cache.get(name)
|
|
if cached is not None:
|
|
return cast(T, cached)
|
|
if not path.exists():
|
|
self._cache[name] = default
|
|
return default
|
|
value = model.model_validate_json(path.read_text(encoding="utf-8"))
|
|
self._cache[name] = value
|
|
return value
|
|
|
|
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)
|
|
self._cache[name] = value
|
|
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)
|