diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 2dc1ff0..ff8b056 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -675,6 +675,18 @@ def _resolve_file_size_bytes( ) +def _extract_archive_size_fields(file_data: Dict[str, Any]) -> Dict[str, Any]: + return { + "sizeKB": file_data.get("sizeKB") or file_data.get("size_kb"), + "sizeBytes": file_data.get("sizeBytes") + or file_data.get("size_bytes") + or file_data.get("fileSize") + or file_data.get("file_size") + or file_data.get("bytes"), + "size": file_data.get("size"), + } + + def _mirror_from_top_file(file_data: Dict[str, Any]) -> Optional[Dict[str, Any]]: url = _normalize_download_url(file_data.get("url") or file_data.get("downloadUrl")) if not url: @@ -691,13 +703,7 @@ def _mirror_from_top_file(file_data: Dict[str, Any]) -> Optional[Dict[str, Any]] "model_id": file_data.get("model_id") or file_data.get("modelId"), "model_version_id": file_data.get("model_version_id") or file_data.get("modelVersionId"), - "sizeKB": file_data.get("sizeKB") or file_data.get("size_kb"), - "sizeBytes": file_data.get("sizeBytes") - or file_data.get("size_bytes") - or file_data.get("fileSize") - or file_data.get("file_size") - or file_data.get("bytes"), - "size": file_data.get("size"), + **_extract_archive_size_fields(file_data), "deletedAt": file_data.get("deletedAt") or file_data.get("deleted_at"), "is_dead": file_data.get("is_dead") or file_data.get("isDead") @@ -743,13 +749,7 @@ def _transform_file_entry(file_data: Dict[str, Any]) -> Dict[str, Any]: "id": file_data.get("id"), "name": name, "type": file_data.get("type"), - "sizeKB": file_data.get("sizeKB") or file_data.get("size_kb"), - "sizeBytes": file_data.get("sizeBytes") - or file_data.get("size_bytes") - or file_data.get("fileSize") - or file_data.get("file_size") - or file_data.get("bytes"), - "size": file_data.get("size"), + **_extract_archive_size_fields(file_data), "downloadUrl": download_url, "primary": bool(file_data.get("primary", file_data.get("is_primary", False))), "mirrors": transformed_mirrors, @@ -785,18 +785,13 @@ def _merge_top_file_mirrors( mirror = _mirror_from_top_file(top_file) if not mirror: continue + size_fields = _extract_archive_size_fields(top_file) merged.append( { "id": top_file.get("id"), "name": mirror.get("filename"), "type": top_file.get("type") or "Model", - "sizeKB": top_file.get("sizeKB") or top_file.get("size_kb"), - "sizeBytes": top_file.get("sizeBytes") - or top_file.get("size_bytes") - or top_file.get("fileSize") - or top_file.get("file_size") - or top_file.get("bytes"), - "size": top_file.get("size"), + **size_fields, "downloadUrl": mirror.get("url"), "primary": True, "mirrors": [mirror], diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index 16020fb..7813306 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -27,9 +27,13 @@ from core.sources.civarchive import ( _build_result_from_normalized_version, _build_result_from_payload, _collect_archive_download_urls, + _extract_archive_size_fields, + _merge_top_file_mirrors, + _mirror_from_top_file, _normalize_archive_version, _normalize_download_url, _resolve_file_size_bytes, + _transform_file_entry, build_civarchive_custom_result, parse_civarchive_url, ) @@ -82,6 +86,45 @@ from core.workflow_updater import ( CIVARCHIVE_BASE = "https://civarchive.com" +class CivArchiveFileSizeAliasTests(unittest.TestCase): + def test_extract_archive_size_fields_preserves_alias_precedence(self): + self.assertEqual( + { + "sizeKB": 12, + "sizeBytes": 34, + "size": 56, + }, + _extract_archive_size_fields( + { + "sizeKB": 0, + "size_kb": 12, + "sizeBytes": 0, + "size_bytes": 0, + "fileSize": 34, + "size": 56, + } + ), + ) + + def test_all_archive_file_projections_use_the_same_size_aliases(self): + file_data = { + "id": 7, + "name": "model.safetensors", + "url": "https://civarchive.com/api/download/models/7", + "size_kb": 12, + "file_size": 34, + } + + mirror = _mirror_from_top_file(file_data) + transformed = _transform_file_entry(file_data) + merged = _merge_top_file_mirrors([], [file_data]) + + for projection in (mirror, transformed, merged[0]): + with self.subTest(projection=projection): + self.assertEqual(12, projection["sizeKB"]) + self.assertEqual(34, projection["sizeBytes"]) + + def write_safetensors_stub(file_path, header): header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8") with open(file_path, "wb") as handle: