Spaces:
Running
Running
File size: 2,248 Bytes
3193174 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import json
from pathlib import Path
from typing import Any
__all__ = ["FileStateStorage", "InMemoryStateStorage"]
class InMemoryStateStorage:
"""Simple in-memory storage for node states."""
def __init__(self) -> None:
self._store: dict[str, dict[str, Any]] = {}
def save(self, node_id: str, state: dict[str, Any]) -> None:
self._store[node_id] = state
def load(self, node_id: str) -> dict[str, Any] | None:
return self._store.get(node_id)
def delete(self, node_id: str) -> None:
self._store.pop(node_id, None)
def keys(self) -> list[str]:
return list(self._store.keys())
def clear(self) -> None:
self._store.clear()
class FileStateStorage:
"""File-based JSON storage for node states."""
def __init__(self, directory: Path | str) -> None:
self._dir = Path(directory)
self._dir.mkdir(parents=True, exist_ok=True)
def _path(self, node_id: str) -> Path:
safe_id = "".join(c if c.isalnum() or c in "-_" else "_" for c in node_id)
return self._dir / f"{safe_id}.json"
def save(self, node_id: str, state: dict[str, Any]) -> None:
path = self._path(node_id)
with path.open("w", encoding="utf-8") as f:
json.dump({"node_id": node_id, "state": state}, f, ensure_ascii=False, indent=2)
def load(self, node_id: str) -> dict[str, Any] | None:
path = self._path(node_id)
if not path.exists():
return None
with path.open("r", encoding="utf-8") as f:
data = json.load(f)
return data.get("state")
def delete(self, node_id: str) -> None:
path = self._path(node_id)
if path.exists():
path.unlink()
def keys(self) -> list[str]:
result = []
for p in self._dir.glob("*.json"):
try:
with p.open("r", encoding="utf-8") as f:
data = json.load(f)
if "node_id" in data:
result.append(data["node_id"])
except (json.JSONDecodeError, KeyError):
continue
return result
def clear(self) -> None:
for p in self._dir.glob("*.json"):
p.unlink()
|