by-aeandCursor Agent
41a54fa376
- Add error handling for JSONDecodeError and UnicodeDecodeError in _read_json - Move corrupt files to timestamped backup (e.g. characters.json.corrupt-1234567890) - Treat corrupt bucket as empty, allowing other buckets to continue working - Clean up temp files if _write_json fails - Log warnings for corrupt file recovery and errors for backup failures Co-authored-by: Cursor Agent <cursoragent@cursor.com>
102 lines
2.7 KiB
Python
102 lines
2.7 KiB
Python
"""File-based character/object/scene storage in ComfyUI user data."""
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import time
|
|
|
|
import folder_paths
|
|
|
|
_STORE_SUBDIR = "byae-ideogram-tools"
|
|
_BUCKETS = frozenset({"characters", "objects", "scenes"})
|
|
_SAFE_NAME = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_\- ]{0,63}$")
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _store_dir() -> str:
|
|
path = os.path.join(folder_paths.get_user_directory(), _STORE_SUBDIR)
|
|
os.makedirs(path, exist_ok=True)
|
|
return path
|
|
|
|
|
|
def _bucket_path(bucket: str) -> str:
|
|
if bucket not in _BUCKETS:
|
|
raise ValueError(f"Unknown store bucket: {bucket}")
|
|
return os.path.join(_store_dir(), f"{bucket}.json")
|
|
|
|
|
|
def _validate_name(name: str) -> str:
|
|
name = (name or "").strip()
|
|
if not _SAFE_NAME.match(name):
|
|
raise ValueError("Invalid name (use letters, numbers, spaces, _ or -)")
|
|
return name
|
|
|
|
|
|
def _read_json(path: str) -> dict:
|
|
if not os.path.isfile(path):
|
|
return {}
|
|
try:
|
|
with open(path, encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
return data if isinstance(data, dict) else {}
|
|
except (json.JSONDecodeError, UnicodeDecodeError) as e:
|
|
# Move corrupt file aside and treat bucket as empty
|
|
timestamp = int(time.time())
|
|
backup_path = f"{path}.corrupt-{timestamp}"
|
|
logger.warning(
|
|
"ByAE Ideogram: corrupt store file %s (%s), moving to %s",
|
|
path, e, backup_path
|
|
)
|
|
try:
|
|
os.rename(path, backup_path)
|
|
except OSError as rename_err:
|
|
logger.error("ByAE Ideogram: failed to backup corrupt file: %s", rename_err)
|
|
return {}
|
|
|
|
|
|
def _write_json(path: str, data: dict) -> None:
|
|
tmp = path + ".tmp"
|
|
try:
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=2, ensure_ascii=False)
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
os.replace(tmp, path)
|
|
except Exception:
|
|
# Clean up temp file on write failure
|
|
if os.path.exists(tmp):
|
|
try:
|
|
os.remove(tmp)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
|
|
|
|
def load_bucket(bucket: str) -> dict:
|
|
return _read_json(_bucket_path(bucket))
|
|
|
|
|
|
def load_all() -> dict:
|
|
return {bucket: load_bucket(bucket) for bucket in _BUCKETS}
|
|
|
|
|
|
def save_entry(bucket: str, name: str, data) -> None:
|
|
name = _validate_name(name)
|
|
path = _bucket_path(bucket)
|
|
store = _read_json(path)
|
|
store[name] = data
|
|
_write_json(path, store)
|
|
|
|
|
|
def delete_entry(bucket: str, name: str) -> bool:
|
|
name = _validate_name(name)
|
|
path = _bucket_path(bucket)
|
|
store = _read_json(path)
|
|
if name not in store:
|
|
return False
|
|
del store[name]
|
|
_write_json(path, store)
|
|
return True
|