Files
novelmaster/noma/scripts/data_modules/snapshot_manager.py
T
2026-06-23 20:29:02 +08:00

93 lines
3.0 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Context snapshot manager.
"""
from __future__ import annotations
import json
from dataclasses import dataclass
from datetime import datetime, timezone
from filelock import FileLock
from pathlib import Path
from typing import Any, Dict, Optional
from .config import get_config
try:
# 当 scripts 目录在 sys.path 中
from security_utils import atomic_write_json
except ImportError: # pragma: no cover
# 当以 python -m scripts.data_modules... 形式运行
from scripts.security_utils import atomic_write_json
SNAPSHOT_VERSION = "1.2"
class SnapshotVersionMismatch(RuntimeError):
def __init__(self, expected: str, actual: str) -> None:
super().__init__(f"snapshot version mismatch: expected {expected}, got {actual}")
self.expected = expected
self.actual = actual
@dataclass
class SnapshotMeta:
chapter: int
version: str
saved_at: str
class SnapshotManager:
def __init__(self, config=None, version: str = SNAPSHOT_VERSION):
self.config = config or get_config()
self.version = version
self.snapshot_dir = self.config.noma_dir / "context_snapshots"
self.snapshot_dir.mkdir(parents=True, exist_ok=True)
def _snapshot_path(self, chapter: int) -> Path:
return self.snapshot_dir / f"ch{chapter:04d}.json"
def _snapshot_lock_path(self, chapter: int) -> Path:
return self._snapshot_path(chapter).with_suffix(".json.lock")
def save_snapshot(self, chapter: int, payload: Dict[str, Any], meta: Optional[Dict[str, Any]] = None) -> Path:
data: Dict[str, Any] = {
"version": self.version,
"chapter": chapter,
"saved_at": datetime.now(timezone.utc).isoformat(),
"payload": payload,
}
if meta:
data["meta"] = meta
path = self._snapshot_path(chapter)
lock = FileLock(str(self._snapshot_lock_path(chapter)), timeout=10)
with lock:
atomic_write_json(path, data, use_lock=False, backup=False)
return path
def load_snapshot(self, chapter: int) -> Optional[Dict[str, Any]]:
path = self._snapshot_path(chapter)
lock = FileLock(str(self._snapshot_lock_path(chapter)), timeout=10)
with lock:
if not path.exists():
return None
data = json.loads(path.read_text(encoding="utf-8"))
version = str(data.get("version", ""))
if version != self.version:
raise SnapshotVersionMismatch(self.version, version)
return data
def delete_snapshot(self, chapter: int) -> bool:
path = self._snapshot_path(chapter)
lock = FileLock(str(self._snapshot_lock_path(chapter)), timeout=10)
with lock:
if path.exists():
path.unlink()
return True
return False
def list_snapshots(self) -> list[str]:
return sorted(p.name for p in self.snapshot_dir.glob("ch*.json"))