wip mlflow implementation

src/state.py

@@ -1,30 +1,65 @@
 import json
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
 STATE_FILE = ".qc-cli.json"
 
 
-def _path(config_dir: str) -> Path:
-    return Path(config_dir) / STATE_FILE
+@dataclass(frozen=True)
+class CliStateStore:
+    config_dir: str = "."
+
+    @property
+    def path(self) -> Path:
+        return Path(self.config_dir) / STATE_FILE
+
+    def read(self) -> dict[str, Any]:
+        if not self.path.exists():
+            return {}
+        with open(self.path) as f:
+            value = json.load(f)
+        return dict(value) if isinstance(value, dict) else {}
+
+    def update(self, **updates: Any) -> None:
+        state = self.read()
+        state.update(updates)
+        self._write(state)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self.read().get(key, default)
+
+    def get_last_training_job(self) -> str | None:
+        value = self.get("last_training_job")
+        return str(value) if value else None
+
+    def set_last_training_job(self, job_name: str) -> None:
+        self.update(last_training_job=job_name)
+
+    def get_training_job(self, job_name: str) -> dict[str, Any]:
+        jobs = self._training_jobs(self.read())
+        value = jobs.get(job_name, {})
+        return dict(value) if isinstance(value, dict) else {}
+
+    def update_training_job(self, job_name: str, **updates: Any) -> None:
+        state = self.read()
+        jobs = self._training_jobs(state)
+        jobs[job_name] = {**jobs.get(job_name, {}), **updates}
+        state["training_jobs"] = jobs
+        self._write(state)
+
+    def set_latest_prerelease_model_version(self, version: str) -> None:
+        self.update(latest_prerelease_model_version=version)
+
+    def _write(self, state: dict[str, Any]) -> None:
+        with open(self.path, "w") as f:
+            json.dump(state, f, indent=2)
+
+    def _training_jobs(self, state: dict[str, Any]) -> dict[str, Any]:
+        value = state.get("training_jobs", {})
+        return dict(value) if isinstance(value, dict) else {}
 
 
-def read_state(config_dir: str = ".") -> dict[str, Any]:
-    path = _path(config_dir)
-    if not path.exists():
-        return {}
-    with open(path) as f:
-        return json.load(f)
-
-
-def write_state(config_dir: str = ".", **updates: str | None) -> None:
-    path = _path(config_dir)
-    state = read_state(config_dir)
-    state.update(updates)
-    with open(path, "w") as f:
-        json.dump(state, f, indent=2)
-
-
-def get_last_training_job(config_dir: str = ".") -> str | None:
-    value = read_state(config_dir).get("last_training_job")
-    return str(value) if value else None
+def store(config_path: str) -> CliStateStore:
+    config_dir = str(Path(config_path).parent)
+    return CliStateStore(config_dir)