Improve init & data handling. Improve typing.

Minor usage change. Instead of setting the hash info directly, call `add_hash` to ensure proper data structure.

modules/hashes.py

@@ -1,5 +1,7 @@
 import hashlib
 import os.path
+from collections import defaultdict
+from typing import TypedDict
 from rich import progress, errors
 from modules.logger import console
 from modules.logger import log
@@ -7,28 +9,37 @@ from modules.json_helpers import readfile, writefile
 from modules.paths import data_path
 
 
-cache_filename = os.path.join(data_path, 'data', 'cache.json')
-cache_data = None
+class HashEntry(TypedDict):
+    mtime: float
+    sha256: str
+
+
+class HashStore(dict[str, HashEntry]):
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def add_hash(self, key: str, mtime: float = 0, sha256: str | None = None):
+        self.__setitem__(key, {"mtime": mtime, "sha256": sha256 or ""})
+
+
+cache_filename = os.path.join(data_path, "data", "cache.json")
+# defaultdict allows for easily using new stores without needing to define them ahead of time
+_data: defaultdict[str, HashStore] = defaultdict(HashStore)
 progress_ok = True
 
 
 def init_cache():
-    global cache_data # pylint: disable=global-statement
-    if cache_data is None:
-        cache_data = {} if not os.path.isfile(cache_filename) else readfile(cache_filename, lock=True, as_type="dict")
+    if os.path.isfile(cache_filename):
+        for store, data in readfile(cache_filename, lock=True, as_type="dict").items():
+            _data[store] = HashStore(data)
 
 
 def dump_cache():
-    writefile(cache_data, cache_filename)
+    writefile(_data, cache_filename)
 
 
-def cache(subsection):
-    global cache_data # pylint: disable=global-statement
-    if cache_data is None:
-        cache_data = {} if not os.path.isfile(cache_filename) else readfile(cache_filename, lock=True, as_type="dict")
-    s = cache_data.get(subsection, {})
-    cache_data[subsection] = s
-    return s
+def cache(store: str = "hashes") -> HashStore:
+    return _data[store]
 
 
 def calculate_sha256(filename, quiet=False):
@@ -55,19 +66,18 @@ def calculate_sha256(filename, quiet=False):
     return hash_sha256.hexdigest()
 
 
-def sha256_from_cache(filename, title, use_addnet_hash=False):
+def sha256_from_cache(filename: str, title: str, use_addnet_hash=False):
     hashes = cache("hashes-addnet") if use_addnet_hash else cache("hashes")
     if title not in hashes:
         return None
-    cached_sha256 = hashes[title].get("sha256", None)
-    cached_mtime = hashes[title].get("mtime", 0)
+    cached = hashes[title]
     ondisk_mtime = os.path.getmtime(filename) if os.path.isfile(filename) else 0
-    if ondisk_mtime > cached_mtime or cached_sha256 is None:
+    if ondisk_mtime > cached["mtime"] or not cached["sha256"]:
         return None
-    return cached_sha256
+    return cached["sha256"]
 
 
-def sha256(filename, title, use_addnet_hash=False):
+def sha256(filename: str, title: str, use_addnet_hash=False):
     from modules import shared
     global progress_ok # pylint: disable=global-statement
     hashes = cache("hashes-addnet") if use_addnet_hash else cache("hashes")
@@ -92,10 +102,7 @@ def sha256(filename, title, use_addnet_hash=False):
                 sha256_value = addnet_hash_safetensors(f)
     else:
         sha256_value = calculate_sha256(filename)
-    hashes[title] = {
-        "mtime": os.path.getmtime(filename),
-        "sha256": sha256_value
-    }
+    hashes.add_hash(title, os.path.getmtime(filename), sha256_value)
     shared.state.end(jobid)
     dump_cache()
     return sha256_value