by-ae-ComfyUI-ByAE-IdeogramTools/store.py

"""File-based character/object/scene storage in ComfyUI user data."""

import contextlib
import json
import logging
import os
import re
import tempfile
import threading
import time
import unicodedata

import folder_paths

_STORE_SUBDIR = "byae-ideogram-tools"
_BUCKETS = frozenset({"characters", "objects", "scenes"})
# Allow Unicode letters, numbers, spaces, underscores, and hyphens
# Must start with a letter or number (Unicode-aware)
_SAFE_NAME = re.compile(r"^[\w][\w\- ]{0,63}$", re.UNICODE)
_MAX_NAME_LEN = 64

logger = logging.getLogger(__name__)

# File locks for concurrent write protection (one lock per bucket).
# The guard lock covers creating the per-bucket lock; without it two threads
# can each install a different lock and then write the same file together.
_BUCKET_LOCKS = {}
_BUCKET_LOCKS_GUARD = threading.Lock()


def _store_dir() -> str:
    path = os.path.join(folder_paths.get_user_directory(), _STORE_SUBDIR)
    os.makedirs(path, exist_ok=True)
    return path


@contextlib.contextmanager
def _bucket_lock(bucket: str):
    """Thread-safe lock for bucket writes to prevent concurrent write corruption."""
    with _BUCKET_LOCKS_GUARD:
        lock = _BUCKET_LOCKS.get(bucket)
        if lock is None:
            lock = threading.Lock()
            _BUCKET_LOCKS[bucket] = lock
    lock.acquire()
    try:
        yield
    finally:
        lock.release()


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:
    # A letter plus a combining accent does not match the composed letter,
    # and the combining mark is not a word character, so the name was rejected.
    name = unicodedata.normalize("NFC", (name or "").strip())
    if len(name) > _MAX_NAME_LEN:
        raise ValueError(f"Invalid name (max {_MAX_NAME_LEN} characters)")
    if not _SAFE_NAME.match(name):
        raise ValueError("Invalid name (use letters, numbers, spaces, _ or -)")
    return name


def _canonicalize_bucket(data: dict) -> dict:
    """One key per composed name. A later entry wins when two keys fold together."""
    cleaned = {}
    for key, value in data.items():
        if not isinstance(key, str):
            continue
        name = unicodedata.normalize("NFC", key.strip())
        if not name:
            continue
        if name in cleaned:
            logger.warning(
                "ByAE Ideogram: store name %r matches preset %r; keeping the later entry",
                key, name,
            )
        cleaned[name] = value
    return cleaned


def _backup_unreadable(path: str, reason) -> None:
    """Move a store file aside. A name that already exists is not replaced."""
    timestamp = int(time.time())
    backup_path = f"{path}.corrupt-{timestamp}"
    if os.path.exists(backup_path):
        backup_path = f"{path}.corrupt-{timestamp}-{time.time_ns()}"
    logger.warning(
        "ByAE Ideogram: corrupt store file %s (%s), moving to %s",
        path, reason, 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)


def _read_json(path: str) -> dict:
    if not os.path.isfile(path):
        return {}
    try:
        # utf-8-sig ignores one leading byte-order mark. utf-8 rejects it, and
        # that used to move an otherwise valid bucket aside as corrupt.
        with open(path, encoding="utf-8-sig") as f:
            data = json.load(f)
    except (json.JSONDecodeError, UnicodeDecodeError) as e:
        _backup_unreadable(path, e)
        return {}
    if isinstance(data, dict):
        return _canonicalize_bucket(data)
    # A JSON array or other value is valid JSON but not a bucket. Leaving it
    # in place made the next save overwrite it with no backup.
    _backup_unreadable(path, f"expected a JSON object, got {type(data).__name__}")
    return {}


def _write_json(path: str, data: dict) -> None:
    # Use a temporary file with a unique name to avoid conflicts during concurrent writes
    temp_fd, tmp = tempfile.mkstemp(
        dir=os.path.dirname(path),
        prefix=".tmp_",
        suffix=".json"
    )
    try:
        with os.fdopen(temp_fd, "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
        try:
            os.close(temp_fd)
        except (OSError, ValueError):
            pass
        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)
    with _bucket_lock(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)
    with _bucket_lock(bucket):
        store = _read_json(path)
        if name not in store:
            return False
        del store[name]
        _write_json(path, store)
        return True