From e326a107f43e1d2f080a34ced4614a13496d9d98 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Wed, 5 Aug 2026 02:13:51 +0200 Subject: [PATCH] refactor: centralize initial download progress --- core/download/aria2_backend.py | 23 ++++++---------- core/download/huggingface_xet.py | 24 +++++++--------- core/download/orchestrator.py | 46 ++++++++++++------------------- core/download/state.py | 27 ++++++++++++++++++ tests/test_download_state.py | 47 ++++++++++++++++++++++++++++++++ 5 files changed, 111 insertions(+), 56 deletions(-) create mode 100644 tests/test_download_state.py diff --git a/core/download/aria2_backend.py b/core/download/aria2_backend.py index fb1e96d..1caa86a 100644 --- a/core/download/aria2_backend.py +++ b/core/download/aria2_backend.py @@ -4,6 +4,7 @@ import socket from typing import Any, Dict, Optional from ..log_system import create_module_logger +from .state import create_initial_progress log = create_module_logger("core.downloader") @@ -724,20 +725,14 @@ def download_file_with_aria2( expected_sha256 = facade._extract_expected_sha256(metadata) with facade.download_lock: - facade.download_progress[download_id] = { - "status": "starting", - "progress": 0, - "total_size": 0, - "downloaded": 0, - "filename": filename, - "path": dest_path, - "directory": facade.os.path.dirname(dest_path), - "url": url, - "error": None, - "speed": 0, - "start_time": start_time, - "download_backend": "aria2", - } + facade.download_progress[download_id] = create_initial_progress( + url=url, + path=dest_path, + filename=filename, + directory=facade.os.path.dirname(dest_path), + download_backend="aria2", + start_time=start_time, + ) try: facade.os.makedirs(facade.os.path.dirname(dest_path), exist_ok=True) diff --git a/core/download/huggingface_xet.py b/core/download/huggingface_xet.py index d3091d2..42a730e 100644 --- a/core/download/huggingface_xet.py +++ b/core/download/huggingface_xet.py @@ -5,6 +5,7 @@ from typing import Any, Dict, Optional from urllib.parse import urlparse from ..log_system import create_module_logger +from .state import create_initial_progress log = create_module_logger("core.downloader") @@ -334,20 +335,15 @@ def download_huggingface_xet( ) with facade.download_lock: - facade.download_progress[download_id] = { - "status": "starting", - "progress": 0, - "total_size": expected_size, - "downloaded": 0, - "filename": filename, - "path": dest_path, - "directory": facade.os.path.dirname(dest_path), - "url": validated_url, - "error": None, - "speed": 0, - "start_time": start_time, - "download_backend": "huggingface_xet", - } + facade.download_progress[download_id] = create_initial_progress( + url=validated_url, + path=dest_path, + filename=filename, + directory=facade.os.path.dirname(dest_path), + download_backend="huggingface_xet", + start_time=start_time, + total_size=expected_size, + ) facade.cancelled_downloads.discard(download_id) try: diff --git a/core/download/orchestrator.py b/core/download/orchestrator.py index 5713918..16f7448 100644 --- a/core/download/orchestrator.py +++ b/core/download/orchestrator.py @@ -2,6 +2,8 @@ from typing import Any, Callable, Dict, Optional +from .state import create_initial_progress + def find_active_download_for_path( download_progress: Dict[str, Dict[str, Any]], @@ -112,20 +114,14 @@ def download_file( last_cli_log = start_time with download_lock: - download_progress[download_id] = { - "status": "starting", - "progress": 0, - "total_size": 0, - "downloaded": 0, - "filename": facade.get_filename_from_path(dest_path), - "path": dest_path, - "directory": os.path.dirname(dest_path), - "url": url, - "error": None, - "speed": 0, - "start_time": start_time, - "download_backend": "python", - } + download_progress[download_id] = create_initial_progress( + url=url, + path=dest_path, + filename=facade.get_filename_from_path(dest_path), + directory=os.path.dirname(dest_path), + download_backend="python", + start_time=start_time, + ) try: destination_directory = os.path.dirname(dest_path) @@ -596,20 +592,14 @@ def start_background_download( f"(download_id={active_download_id})" ) return active_download_id - download_progress[download_id] = { - "status": "starting", - "progress": 0, - "total_size": 0, - "downloaded": 0, - "filename": filename, - "path": initial_path, - "directory": initial_directory, - "url": url, - "error": None, - "speed": 0, - "start_time": time.time(), - "download_backend": facade._download_backend_from_settings(), - } + download_progress[download_id] = create_initial_progress( + url=url, + path=initial_path, + filename=filename, + directory=initial_directory, + download_backend=facade._download_backend_from_settings(), + start_time=time.time(), + ) def run_download() -> None: try: diff --git a/core/download/state.py b/core/download/state.py index df3d18a..610139b 100644 --- a/core/download/state.py +++ b/core/download/state.py @@ -16,6 +16,33 @@ xet_transfers: Dict[str, Dict[str, Any]] = {} xet_transfers_lock = threading.Lock() +def create_initial_progress( + *, + url: str, + path: str, + filename: str, + directory: str, + download_backend: str, + start_time: float, + total_size: int = 0, +) -> Dict[str, Any]: + """Build the shared initial progress record for a download.""" + return { + "status": "starting", + "progress": 0, + "total_size": total_size, + "downloaded": 0, + "filename": filename, + "path": path, + "directory": directory, + "url": url, + "error": None, + "speed": 0, + "start_time": start_time, + "download_backend": download_backend, + } + + @dataclass(frozen=True) class DownloadStateDependencies: """State and side-effect dependencies supplied by the composition layer.""" diff --git a/tests/test_download_state.py b/tests/test_download_state.py new file mode 100644 index 0000000..28a4e0e --- /dev/null +++ b/tests/test_download_state.py @@ -0,0 +1,47 @@ +import importlib + + +def test_initial_download_progress_preserves_the_shared_contract(): + state = importlib.import_module("core.download.state") + create_initial_progress = getattr(state, "create_initial_progress", None) + assert callable(create_initial_progress) + + progress = create_initial_progress( + url="https://example.com/model.safetensors", + path=r"C:\models\model.safetensors", + filename="model.safetensors", + directory=r"C:\models", + download_backend="aria2", + total_size=42, + start_time=12.5, + ) + + assert progress == { + "status": "starting", + "progress": 0, + "total_size": 42, + "downloaded": 0, + "filename": "model.safetensors", + "path": r"C:\models\model.safetensors", + "directory": r"C:\models", + "url": "https://example.com/model.safetensors", + "error": None, + "speed": 0, + "start_time": 12.5, + "download_backend": "aria2", + } + + +def test_initial_download_progress_defaults_total_size_to_zero(): + state = importlib.import_module("core.download.state") + progress = state.create_initial_progress( + url="https://example.com/model.safetensors", + path="model.safetensors", + filename="model.safetensors", + directory="", + download_backend="python", + start_time=1.0, + ) + + assert progress["total_size"] == 0 + assert progress["download_backend"] == "python"