From 3d98572a62d5e97a3095e4d0b8a98e1fcada99b5 Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Tue, 14 Oct 2025 16:03:33 +0800
Subject: [PATCH] feat: improve civitai data handling and type safety, fixes
 #565

- Replace setdefault with get and explicit dict initialization in MetadataUpdater
- Change civitai field type from Optional[Dict] to Dict[str, Any] with default_factory
- Add None check and dict initialization in BaseModelMetadata.__post_init__
- Ensures civitai data is always a dictionary, preventing type errors and improving code reliability
---
 py/utils/example_images_metadata.py   |  9 +++++--
 py/utils/models.py                    |  5 +++-
 tests/utils/test_metadata_defaults.py | 37 +++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 3 deletions(-)
 create mode 100644 tests/utils/test_metadata_defaults.py

diff --git a/py/utils/example_images_metadata.py b/py/utils/example_images_metadata.py
index 5772e573..72ddb951 100644
--- a/py/utils/example_images_metadata.py
+++ b/py/utils/example_images_metadata.py
@@ -270,7 +270,12 @@ class MetadataUpdater:
         """
         try:
             await MetadataManager.hydrate_model_data(model_data)
-            civitai_data = model_data.setdefault('civitai', {})
+            civitai_data = model_data.get('civitai')
+
+            if not isinstance(civitai_data, dict):
+                civitai_data = {}
+                model_data['civitai'] = civitai_data
+
             custom_images = civitai_data.get('customImages')
 
             if not isinstance(custom_images, list):
@@ -446,4 +451,4 @@ class MetadataUpdater:
             
         except Exception as e:
             logger.error(f"Error parsing image metadata: {e}", exc_info=True)
-            return None
\ No newline at end of file
+            return None
diff --git a/py/utils/models.py b/py/utils/models.py
index 4caffa3e..75acb840 100644
--- a/py/utils/models.py
+++ b/py/utils/models.py
@@ -18,7 +18,7 @@ class BaseModelMetadata:
     preview_nsfw_level: int = 0 # NSFW level of the preview image
     notes: str = ""             # Additional notes
     from_civitai: bool = True   # Whether from Civitai
-    civitai: Optional[Dict] = None  # Civitai API data if available
+    civitai: Dict[str, Any] = field(default_factory=dict)  # Civitai API data if available
     tags: List[str] = None      # Model tags
     modelDescription: str = ""  # Full model description
     civitai_deleted: bool = False  # Whether deleted from Civitai
@@ -31,6 +31,9 @@ class BaseModelMetadata:
 
     def __post_init__(self):
         # Initialize empty lists to avoid mutable default parameter issue
+        if self.civitai is None:
+            self.civitai = {}
+
         if self.tags is None:
             self.tags = []
 
diff --git a/tests/utils/test_metadata_defaults.py b/tests/utils/test_metadata_defaults.py
new file mode 100644
index 00000000..94748354
--- /dev/null
+++ b/tests/utils/test_metadata_defaults.py
@@ -0,0 +1,37 @@
+import json
+import pytest
+
+from py.utils.metadata_manager import MetadataManager
+from py.utils.models import BaseModelMetadata
+
+
+@pytest.mark.asyncio
+async def test_base_model_metadata_sets_empty_civitai_dict():
+    metadata = BaseModelMetadata(
+        file_name="model",
+        model_name="Model",
+        file_path="/tmp/model.safetensors",
+        size=0,
+        modified=0.0,
+        sha256="deadbeef",
+        base_model="Unknown",
+        preview_url="",
+    )
+
+    assert metadata.civitai == {}
+
+
+@pytest.mark.asyncio
+async def test_create_default_metadata_uses_empty_civitai(tmp_path):
+    model_path = tmp_path / "example.safetensors"
+    model_path.write_bytes(b"stub")
+
+    metadata = await MetadataManager.create_default_metadata(str(model_path))
+
+    assert metadata is not None
+    assert metadata.civitai == {}
+
+    metadata_path = model_path.with_suffix(".metadata.json")
+    payload = json.loads(metadata_path.read_text(encoding="utf-8"))
+
+    assert payload.get("civitai") == {}