diff --git a/core/metadata_audit.py b/core/metadata_audit.py index 1f42439..1dadc41 100644 --- a/core/metadata_audit.py +++ b/core/metadata_audit.py @@ -7,7 +7,7 @@ Provides audits for local model sidecar metadata files. import math import os from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, Iterable, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple from .log_system import create_module_logger from .metadata_model_utils import dedupe_models, is_model_file_path @@ -18,7 +18,7 @@ from .path_utils import ( get_path_identity, read_merged_model_metadata, ) -from .type_utils import FILE_SIZE_BYTE_KEYS, FILE_SIZE_KIB_KEYS, format_size_bytes +from .type_utils import extract_file_size_detail, format_size_bytes from .worker_utils import resolve_worker_count log = create_module_logger(__name__) @@ -56,46 +56,15 @@ def _empty_audit_counts() -> Dict[str, Any]: } -def _coerce_size_bytes(value: Any, multiplier: int = 1) -> Optional[int]: - if value is None or value == "" or isinstance(value, bool): - return None - - try: - if isinstance(value, str): - text = value.strip().replace(",", "").replace("_", "") - if not text or text.lower() in {"none", "null", "undefined"}: - return None - number = float(text) - else: - number = float(value) - except (TypeError, ValueError): - return None - - if not math.isfinite(number) or number < 0: - return None - - return int(number * multiplier) - - -def _iter_size_fields(data: Dict[str, Any], prefix: str = "") -> Iterable[Tuple[str, Any, int]]: - for key in FILE_SIZE_BYTE_KEYS: - if key in data: - yield (f"{prefix}{key}", data.get(key), 1) - for key in FILE_SIZE_KIB_KEYS: - if key in data: - yield (f"{prefix}{key}", data.get(key), 1024) - - def _extract_size_from_object(data: Any, prefix: str = "") -> Optional[Tuple[int, str]]: if not isinstance(data, dict): return None - from .type_utils import extract_file_size - size = extract_file_size(data) - if size is not None: - for field_path, value, multiplier in _iter_size_fields(data, prefix): - if _coerce_size_bytes(value, multiplier=multiplier) == size: - return size, field_path + detail = extract_file_size_detail(data) + if detail is not None: + size, field_name = detail + if field_name: + return size, f"{prefix}{field_name}" return size, f"{prefix}size" return None diff --git a/core/type_utils.py b/core/type_utils.py index 326f6de..3373e56 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -4,6 +4,7 @@ Type Utilities Module Unified helper functions for safe data type casting and normalization. """ +import math import re from typing import Any, Dict, List @@ -1116,10 +1117,28 @@ def clear_remote_size_cache() -> None: _remote_size_cache.clear() -def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]: - """ - Extracts file size in bytes from various metadata formats (HuggingFace, CivitAI, CivArchive). - """ +def _coerce_size_field(value: Any, multiplier: int = 1) -> Optional[int]: + if value is None or value == "" or isinstance(value, bool): + return None + + try: + if isinstance(value, str): + text = value.strip().replace(",", "").replace("_", "") + if not text or text.lower() in {"none", "null", "undefined"}: + return None + number = float(text) + else: + number = float(value) + except (TypeError, ValueError): + return None + + if not math.isfinite(number) or number < 0: + return None + + return int(number * multiplier) + + +def _extract_file_size_value(file_info: Dict[str, Any]) -> Optional[int]: if not isinstance(file_info, dict): return None @@ -1155,13 +1174,38 @@ def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]: mirrors = [mirrors] for mirror in mirrors: if isinstance(mirror, dict): - size = extract_file_size(mirror) + size = _extract_file_size_value(mirror) if size: return size return None +def extract_file_size_detail( + file_info: Dict[str, Any], +) -> Optional[Tuple[int, Optional[str]]]: + """Extract a file size and its matching direct metadata field, when available.""" + size = _extract_file_size_value(file_info) + if size is None: + return None + + if isinstance(file_info, dict): + for key in FILE_SIZE_BYTE_KEYS: + if _coerce_size_field(file_info.get(key)) == size: + return size, key + for key in FILE_SIZE_KIB_KEYS: + if _coerce_size_field(file_info.get(key), multiplier=1024) == size: + return size, key + + return size, None + + +def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]: + """Extract a file size in bytes from supported provider metadata formats.""" + detail = extract_file_size_detail(file_info) + return detail[0] if detail else None + + def normalize_category_to_model_type(category: str) -> str: """ Map a download or ComfyUI category name, such as "checkpoints", to a diff --git a/tests/test_file_size_aliases.py b/tests/test_file_size_aliases.py index 95cda3a..15d4025 100644 --- a/tests/test_file_size_aliases.py +++ b/tests/test_file_size_aliases.py @@ -1,7 +1,7 @@ import unittest from core.metadata_audit import extract_metadata_size -from core.type_utils import extract_file_size +from core.type_utils import extract_file_size, extract_file_size_detail class FileSizeAliasTests(unittest.TestCase): @@ -29,8 +29,23 @@ class FileSizeAliasTests(unittest.TestCase): metadata = {"size": 1024, "sizeBytes": 2048} self.assertEqual(2048, extract_file_size(metadata)) + self.assertEqual((2048, "sizeBytes"), extract_file_size_detail(metadata)) self.assertEqual((2048, "sizeBytes"), extract_metadata_size(metadata, "model.safetensors")) + def test_detail_keeps_legacy_direct_field_matching_order(self): + self.assertEqual( + (2048, "size"), + extract_file_size_detail({"size": 2048, "sizeBytes": 2048}), + ) + self.assertEqual( + (2048, "sizeBytes"), + extract_file_size_detail({"sizeKB": 2, "sizeBytes": 2048}), + ) + self.assertEqual( + (300, None), + extract_file_size_detail({"lfs": {"size": 300}}), + ) + if __name__ == "__main__": unittest.main()