refactor: centralize file size aliases

core/metadata_audit.py

@@ -17,14 +17,12 @@ from .path_utils import (
     get_path_identity,
     read_merged_model_metadata,
 )
-from .type_utils import format_size_bytes
+from .type_utils import FILE_SIZE_BYTE_KEYS, FILE_SIZE_KIB_KEYS, format_size_bytes
 from .worker_utils import resolve_worker_count
 
 log = create_module_logger(__name__)
 
 
-BYTE_SIZE_KEYS = ("size", "sizeBytes", "size_bytes", "fileSize", "file_size", "bytes")
-KIB_SIZE_KEYS = ("sizeKB", "size_kb")
 MAX_AUDIT_WORKERS = 64
 MIN_AUDIT_WORKERS = 1
 MIN_AUDIT_BATCH_SIZE = 16
@@ -68,10 +66,10 @@ def _coerce_size_bytes(value: Any, multiplier: int = 1) -> Optional[int]:
 
 
 def _iter_size_fields(data: Dict[str, Any], prefix: str = "") -> Iterable[Tuple[str, Any, int]]:
-    for key in BYTE_SIZE_KEYS:
+    for key in FILE_SIZE_BYTE_KEYS:
         if key in data:
             yield (f"{prefix}{key}", data.get(key), 1)
-    for key in KIB_SIZE_KEYS:
+    for key in FILE_SIZE_KIB_KEYS:
         if key in data:
             yield (f"{prefix}{key}", data.get(key), 1024)
 

core/type_utils.py

@@ -118,6 +118,11 @@ QUANTIZATION_LEVEL_SUFFIXES = (
 )
 
 
+FILE_SIZE_BYTE_KEYS = ("size", "sizeBytes", "size_bytes", "fileSize", "file_size", "bytes")
+FILE_SIZE_KIB_KEYS = ("sizeKB", "size_kb")
+FILE_SIZE_LFS_KEYS = FILE_SIZE_BYTE_KEYS[:3]
+
+
 QUANTIZATION_SCHEME_SUFFIXES = (
     "gptq",
     "awq",
@@ -1119,7 +1124,7 @@ def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]:
         return None
 
     # 1. Check direct sizeKB or size_kb keys (e.g. from CivArchive or CivitAI)
-    for key in ("sizeKB", "size_kb"):
+    for key in FILE_SIZE_KIB_KEYS:
         val = file_info.get(key)
         if val is not None and val != "":
             try:
@@ -1128,7 +1133,7 @@ def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]:
                 pass
 
     # 2. Check direct size bytes keys
-    for key in ("sizeBytes", "size_bytes", "fileSize", "file_size", "bytes", "size"):
+    for key in (*FILE_SIZE_BYTE_KEYS[1:], FILE_SIZE_BYTE_KEYS[0]):
         val = file_info.get(key)
         if val is not None and val != "":
             size = parse_size_header(val)
@@ -1138,7 +1143,7 @@ def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]:
     # 3. Check nested LFS info (common in HuggingFace API metadata)
     lfs_info = file_info.get("lfs")
     if isinstance(lfs_info, dict):
-        for key in ("size", "sizeBytes", "size_bytes"):
+        for key in FILE_SIZE_LFS_KEYS:
             size = parse_size_header(lfs_info.get(key))
             if size:
                 return size

tests/test_file_size_aliases.py

@@ -0,0 +1,36 @@
+import unittest
+
+from core.metadata_audit import extract_metadata_size
+from core.type_utils import extract_file_size
+
+
+class FileSizeAliasTests(unittest.TestCase):
+
+    def test_byte_aliases_are_supported_by_both_extractors(self):
+        byte_aliases = ("size", "sizeBytes", "size_bytes", "fileSize", "file_size", "bytes")
+
+        for key in byte_aliases:
+            with self.subTest(key=key):
+                self.assertEqual(2048, extract_file_size({key: 2048}))
+                self.assertEqual((2048, key), extract_metadata_size({key: 2048}, "model.safetensors"))
+
+    def test_kibibyte_aliases_are_supported_by_both_extractors(self):
+        for key in ("sizeKB", "size_kb"):
+            with self.subTest(key=key):
+                self.assertEqual(2048, extract_file_size({key: 2}))
+                self.assertEqual((2048, key), extract_metadata_size({key: 2}, "model.safetensors"))
+
+    def test_lfs_byte_aliases_remain_supported(self):
+        for key in ("size", "sizeBytes", "size_bytes"):
+            with self.subTest(key=key):
+                self.assertEqual(2048, extract_file_size({"lfs": {key: 2048}}))
+
+    def test_existing_extractor_precedence_is_preserved_when_aliases_collide(self):
+        metadata = {"size": 1024, "sizeBytes": 2048}
+
+        self.assertEqual(2048, extract_file_size(metadata))
+        self.assertEqual((2048, "sizeBytes"), extract_metadata_size(metadata, "model.safetensors"))
+
+
+if __name__ == "__main__":
+    unittest.main()