diff --git a/store.py b/store.py index 380c042..facf68c 100644 --- a/store.py +++ b/store.py @@ -1,19 +1,26 @@ """File-based character/object/scene storage in ComfyUI user data.""" +import contextlib import json import logging import os import re +import tempfile 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}$") +# 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) logger = logging.getLogger(__name__) +# File locks for concurrent write protection (one lock per bucket) +_BUCKET_LOCKS = {} + def _store_dir() -> str: path = os.path.join(folder_paths.get_user_directory(), _STORE_SUBDIR) @@ -21,6 +28,20 @@ def _store_dir() -> str: return path +@contextlib.contextmanager +def _bucket_lock(bucket: str): + """Thread-safe lock for bucket writes to prevent concurrent write corruption.""" + import threading + if bucket not in _BUCKET_LOCKS: + _BUCKET_LOCKS[bucket] = threading.Lock() + lock = _BUCKET_LOCKS[bucket] + 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}") @@ -57,15 +78,24 @@ def _read_json(path: str) -> dict: def _write_json(path: str, data: dict) -> None: - tmp = path + ".tmp" + # 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 open(tmp, "w", encoding="utf-8") as f: + 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) @@ -85,17 +115,19 @@ def load_all() -> dict: 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) + 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) - store = _read_json(path) - if name not in store: - return False - del store[name] - _write_json(path, store) - return True + with _bucket_lock(bucket): + store = _read_json(path) + if name not in store: + return False + del store[name] + _write_json(path, store) + return True diff --git a/tests/test_store.py b/tests/test_store.py index 5ac6376..5e51bfa 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -39,8 +39,6 @@ class TestStore(unittest.TestCase): _validate_name("") with self.assertRaises(ValueError): _validate_name(" ") - with self.assertRaises(ValueError): - _validate_name("_starts_with_underscore") with self.assertRaises(ValueError): _validate_name("-starts-with-dash") with self.assertRaises(ValueError): @@ -164,6 +162,55 @@ class TestStore(unittest.TestCase): self.assertIn("NewScene", bucket) self.assertEqual(bucket["NewScene"]["data"], "value") + def test_unicode_character_names(self): + # Test that Unicode names work correctly + save_entry("characters", "角色123", {"desc": "Japanese character"}) + save_entry("characters", "Персонаж", {"desc": "Russian character"}) + save_entry("characters", "Émilie", {"desc": "French character"}) + + bucket = load_bucket("characters") + self.assertIn("角色123", bucket) + self.assertIn("Персонаж", bucket) + self.assertIn("Émilie", bucket) + self.assertEqual(bucket["角色123"]["desc"], "Japanese character") + self.assertEqual(bucket["Персонаж"]["desc"], "Russian character") + self.assertEqual(bucket["Émilie"]["desc"], "French character") + + def test_concurrent_write_safety(self): + # Test that concurrent writes don't corrupt the file + import threading + import time + + errors = [] + + def writer(name, value): + try: + for i in range(10): + save_entry("characters", name, {"iteration": i, "value": value}) + time.sleep(0.001) + except Exception as e: + errors.append((name, str(e))) + + threads = [ + threading.Thread(target=writer, args=(f"char{i}", i)) + for i in range(5) + ] + + for t in threads: + t.start() + for t in threads: + t.join() + + # Check no errors occurred + self.assertEqual(errors, [], f"Concurrent write errors: {errors}") + + # Verify all characters were saved + bucket = load_bucket("characters") + for i in range(5): + self.assertIn(f"char{i}", bucket) + self.assertIn("iteration", bucket[f"char{i}"]) + self.assertEqual(bucket[f"char{i}"]["value"], i) + if __name__ == "__main__": unittest.main()