From cb666902e5bb17fdc0fa0f0cada5f3be01cac7fa Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sat, 4 Jul 2026 00:18:41 +0200 Subject: [PATCH] refactor: centralize hash calculation callbacks and catalog backup logic --- __init__.py | 92 +++++++++++++++++++------------------- core/path_utils.py | 46 +++++++++++++++++-- core/sources/model_list.py | 8 +--- core/sources/popular.py | 16 +------ 4 files changed, 93 insertions(+), 69 deletions(-) diff --git a/__init__.py b/__init__.py index 5dd5b6b..e37da71 100644 --- a/__init__.py +++ b/__init__.py @@ -724,8 +724,7 @@ class ModelResolverExtension: # cleanup_hash_progress, update_hash_progress, is_hash_progress_cancelled, and mark_hash_progress_cancelled removed - class HashCalculationCancelled(Exception): - pass + from .core.path_utils import HashCalculationCancelled def resolve_hash_file_request(data): import os as _os @@ -799,14 +798,10 @@ class ModelResolverExtension: return resolved_metadata_path, metadata_updated def calculate_sha256_with_progress(normalized_path, progress_id=""): - import hashlib import os as _os + from .core.path_utils import calculate_file_sha256 as _calculate_file_sha256_core total_bytes = max(0, _os.path.getsize(normalized_path)) - bytes_read = 0 - sha256_hash = hashlib.sha256() - last_update = 0.0 - chunk_size = 1024 * 1024 * 4 self.hash_tracker.update( progress_id, @@ -818,45 +813,52 @@ class ModelResolverExtension: total_bytes=total_bytes, ) - with open(normalized_path, "rb") as handle: - for chunk in iter(lambda: handle.read(chunk_size), b""): - if not chunk: - continue - sha256_hash.update(chunk) - bytes_read += len(chunk) - if self.hash_tracker.is_cancelled(progress_id): - percent = 0 if total_bytes <= 0 else min( - 98, - (bytes_read / total_bytes) * 98, - ) - self.hash_tracker.update( - progress_id, - status="cancelled", - stage="cancelled", - message="Hash calculation cancelled", - percent=percent, - bytes_read=bytes_read, - total_bytes=total_bytes, - ) - raise HashCalculationCancelled() - now = time.time() - if now - last_update >= 0.15 or bytes_read >= total_bytes: - percent = 98 if total_bytes <= 0 else min( - 98, - (bytes_read / total_bytes) * 98, - ) - self.hash_tracker.update( - progress_id, - status="running", - stage="hashing", - message="Calculating SHA256...", - percent=percent, - bytes_read=bytes_read, - total_bytes=total_bytes, - ) - last_update = now + _bytes_read = [0] + _last_update = [0.0] - return sha256_hash.hexdigest() + def on_progress(bytes_read, total_bytes): + _bytes_read[0] = bytes_read + now = time.time() + if now - _last_update[0] >= 0.15 or bytes_read >= total_bytes: + percent = 98 if total_bytes <= 0 else min( + 98, + (bytes_read / total_bytes) * 98, + ) + self.hash_tracker.update( + progress_id, + status="running", + stage="hashing", + message="Calculating SHA256...", + percent=percent, + bytes_read=bytes_read, + total_bytes=total_bytes, + ) + _last_update[0] = now + + def is_cancelled(): + if self.hash_tracker.is_cancelled(progress_id): + percent = 0 if total_bytes <= 0 else min( + 98, + (_bytes_read[0] / total_bytes) * 98, + ) + self.hash_tracker.update( + progress_id, + status="cancelled", + stage="cancelled", + message="Hash calculation cancelled", + percent=percent, + bytes_read=_bytes_read[0], + total_bytes=total_bytes, + ) + return True + return False + + return _calculate_file_sha256_core( + normalized_path, + chunk_size=1024 * 1024 * 4, + on_progress=on_progress, + is_cancelled=is_cancelled, + ) @routes.post("/model_resolver/calculate-file-hash") @json_api_endpoint("calculate-file-hash") diff --git a/core/path_utils.py b/core/path_utils.py index c829d97..9ab0255 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -8,7 +8,7 @@ and directory matching to be shared across all modules. import os import json import hashlib -from typing import Any, Tuple, Optional +from typing import Any, Tuple, Optional, Callable _log = None @@ -178,22 +178,62 @@ METADATA_DIR = os.path.join( ) -def calculate_file_sha256(file_path: str, chunk_size: int = 131072) -> Optional[str]: - """Calculate SHA256 hex digest for an existing local file safely.""" +class HashCalculationCancelled(Exception): + """Exception raised when file SHA256 calculation is cancelled by the caller.""" + pass + + +def calculate_file_sha256( + file_path: str, + chunk_size: int = 131072, + on_progress: Optional[Callable[[int, int], None]] = None, + is_cancelled: Optional[Callable[[], bool]] = None, +) -> Optional[str]: + """Calculate SHA256 hex digest for an existing local file safely, supporting progress callbacks and cancellation check.""" if not file_path or not os.path.exists(file_path): return None sha256_hash = hashlib.sha256() try: + total_bytes = os.path.getsize(file_path) if on_progress else 0 + bytes_read = 0 with open(file_path, "rb") as f: for byte_block in iter(lambda: f.read(chunk_size), b""): + if is_cancelled and is_cancelled(): + raise HashCalculationCancelled("Hash calculation cancelled") if byte_block: sha256_hash.update(byte_block) + bytes_read += len(byte_block) + if on_progress: + on_progress(bytes_read, total_bytes) return sha256_hash.hexdigest() + except HashCalculationCancelled: + raise except Exception as e: _get_log().error(f"Error computing hash for {file_path}: {e}") return None +def save_catalog_with_backup( + data_path: str, data: Any, meta_path: str, meta: Any, indent: int = 2 +) -> None: + """Backup an existing catalog file to .bak, then atomically write the new data and its metadata.""" + import shutil + # Ensure directory exists + dir_name = os.path.dirname(data_path) + if dir_name: + os.makedirs(dir_name, exist_ok=True) + + # Backup existing data file + if os.path.exists(data_path): + shutil.copy2(data_path, f"{data_path}.bak") + + # Write data file atomically + write_json_atomic(data_path, data, indent=indent) + + # Write metadata file atomically + write_json_atomic(meta_path, meta, indent=indent) + + def read_json_safe(file_path: str, default: Any = None) -> Any: """Safely read and parse a JSON file, returning a default value on error.""" if not file_path or not os.path.exists(file_path): diff --git a/core/sources/model_list.py b/core/sources/model_list.py index c92b4ef..509478f 100644 --- a/core/sources/model_list.py +++ b/core/sources/model_list.py @@ -19,7 +19,7 @@ from ..log_system import create_module_logger log = create_module_logger(__name__) -from ..path_utils import METADATA_DIR, read_json_safe, write_json_atomic +from ..path_utils import METADATA_DIR, read_json_safe, write_json_atomic, save_catalog_with_backup MODEL_LIST_FILE = os.path.join(METADATA_DIR, "model-list.json") MODEL_LIST_META_FILE = os.path.join(METADATA_DIR, "model-list.meta.json") MODEL_LIST_SOURCE_URL = ( @@ -152,10 +152,6 @@ def update_model_list_from_remote() -> Dict[str, Any]: if not isinstance(models, list) or not models: raise ValueError("Downloaded model-list.json does not contain a non-empty models list") - if os.path.exists(MODEL_LIST_FILE): - shutil.copy2(MODEL_LIST_FILE, f"{MODEL_LIST_FILE}.bak") - - write_json_atomic(MODEL_LIST_FILE, data, indent=2) meta = { "source": "Comfy-Org/ComfyUI-Manager", "source_url": MODEL_LIST_SOURCE_URL, @@ -166,7 +162,7 @@ def update_model_list_from_remote() -> Dict[str, Any]: "model_count": len(models), "updated_at": _utc_now_iso(), } - write_json_atomic(MODEL_LIST_META_FILE, meta, indent=2) + save_catalog_with_backup(MODEL_LIST_FILE, data, MODEL_LIST_META_FILE, meta, indent=2) reload_model_list() return { diff --git a/core/sources/popular.py b/core/sources/popular.py index 69c787a..a2984c9 100644 --- a/core/sources/popular.py +++ b/core/sources/popular.py @@ -349,17 +349,6 @@ def update_base_models_from_remote() -> Dict[str, Any]: new_added_count += 1 new_added_names.append(name) - # Write updated models to base-models.json atomically - os.makedirs(METADATA_DIR, exist_ok=True) - - # Backup existing - if os.path.exists(BASE_MODELS_FILE): - shutil.copy2(BASE_MODELS_FILE, f"{BASE_MODELS_FILE}.bak") - - # Write atomically - write_json_atomic(BASE_MODELS_FILE, {"base_models": updated_models}, indent=2) - - # Write meta file now_str = datetime.now(timezone.utc).replace(microsecond=0).isoformat() meta = { "updated_at": now_str, @@ -368,10 +357,7 @@ def update_base_models_from_remote() -> Dict[str, Any]: "new_models_added": new_added_count, "new_models_added_list": new_added_names } - - write_json_atomic(BASE_MODELS_META_FILE, meta, indent=2) - - # Force reload database in-memory cache + save_catalog_with_backup(BASE_MODELS_FILE, {"base_models": updated_models}, BASE_MODELS_META_FILE, meta, indent=2) reload_databases() return {