diff --git a/core/download/aria2_backend.py b/core/download/aria2_backend.py index b18651a..e8ac97d 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 .dependencies import require_download_dependencies from .state import create_initial_progress log = create_module_logger("core.downloader") @@ -13,16 +14,9 @@ class Aria2Error(RuntimeError): """Raised when the aria2 backend cannot start or process a request.""" -def _require_dependencies(dependencies: Any) -> Any: - """Return explicitly supplied services for the aria2 backend.""" - if dependencies is None: - raise RuntimeError("aria2 backend dependencies were not provided") - return dependencies - - def try_certifi_ca_path(*, dependencies: Any = None) -> str: """Return certifi's CA bundle path when it is available.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") try: import certifi # type: ignore @@ -38,7 +32,7 @@ def resolve_aria2c_executable( dependencies: Any = None, ) -> str: """Resolve aria2c while restricting explicit paths to the managed install.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") active_settings = ( settings if isinstance(settings, dict) else facade.load_settings() ) @@ -121,7 +115,7 @@ def delete_partial_download_files( dependencies: Any = None, ) -> None: """Delete an incomplete model and its aria2 control sidecar.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") for path in (dest_path, f"{dest_path}.aria2"): try: if path and facade.os.path.exists(path): @@ -136,7 +130,7 @@ def delete_python_partial_download_file( dependencies: Any = None, ) -> None: """Remove a partial Python download without touching the final model path.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") try: if partial_path and facade.os.path.exists(partial_path): facade.os.remove(partial_path) @@ -153,7 +147,7 @@ def delete_xet_partial_file( dependencies: Any = None, ) -> bool: """Delete a stopped Xet partial file, retrying while Windows releases it.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") attempts = max(1, int(attempts or 1)) last_error: Optional[Exception] = None for attempt in range(attempts): @@ -188,7 +182,7 @@ def resolve_download_url_for_aria2( dependencies: Any = None, ) -> tuple[str, Dict[str, str]]: """Preflight an aria2 URL and validate every redirect before RPC handoff.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") request_headers = facade.build_download_headers(url, headers) source_host = facade.urlparse(str(url or "")).hostname is_huggingface_source = facade.host_matches_domain( @@ -221,7 +215,7 @@ def resolve_download_url_for_aria2( def read_aria2_version(executable: str, *, dependencies: Any = None) -> str: """Read the installed aria2 version without raising process errors.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") if not executable: return "" try: @@ -265,7 +259,7 @@ def get_aria2_status( dependencies: Any = None, ) -> Dict[str, Any]: """Return aria2 availability, daemon and active-transfer information.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") active_settings = ( settings if isinstance(settings, dict) else facade.load_settings() ) @@ -317,7 +311,7 @@ def aria2_rpc( dependencies: Any = None, ) -> Any: """Call the local aria2 JSON-RPC endpoint.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") if not facade.aria2_rpc_url: raise Aria2Error("aria2 RPC endpoint is not initialized") @@ -356,7 +350,7 @@ def aria2_rpc( def aria2_ping(*, dependencies: Any = None) -> bool: """Check whether the local aria2 RPC endpoint responds.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") try: result = facade._aria2_rpc( "aria2.getVersion", @@ -369,7 +363,7 @@ def aria2_ping(*, dependencies: Any = None) -> bool: def cancel_aria2_idle_timer_locked(*, dependencies: Any = None) -> None: """Cancel the pending idle timer while holding the aria2 lock.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") if facade.aria2_idle_timer is not None: facade.aria2_idle_timer.cancel() facade.aria2_idle_timer = None @@ -377,7 +371,7 @@ def cancel_aria2_idle_timer_locked(*, dependencies: Any = None) -> None: def aria2_has_active_transfers_locked(*, dependencies: Any = None) -> bool: """Return whether aria2 has active resolver-owned transfers.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") return bool(facade.aria2_transfers) @@ -387,7 +381,7 @@ def stop_aria2_daemon( dependencies: Any = None, ) -> Dict[str, Any]: """Stop the aria2 RPC process started by Model Resolver.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") with facade.aria2_lock: facade._cancel_aria2_idle_timer_locked() process = facade.aria2_process @@ -442,7 +436,7 @@ def start_aria2_daemon( dependencies: Any = None, ) -> Dict[str, Any]: """Start the aria2 RPC process without creating a download.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") active_settings = ( settings if isinstance(settings, dict) else facade.load_settings() ) @@ -476,7 +470,7 @@ def start_aria2_daemon( def aria2_idle_stop_worker(*, dependencies: Any = None) -> None: """Stop a resolver-owned daemon after its idle timeout.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") with facade.aria2_lock: facade.aria2_idle_timer = None @@ -497,7 +491,7 @@ def aria2_idle_stop_worker(*, dependencies: Any = None) -> None: def schedule_aria2_idle_stop(*, dependencies: Any = None) -> None: """Schedule daemon shutdown when no resolver transfer remains active.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") settings = facade.load_settings() if not settings.get("aria2_auto_stop_daemon", True): return @@ -525,7 +519,7 @@ def ensure_aria2_daemon( dependencies: Any = None, ) -> None: """Start or reuse the resolver-owned aria2 daemon.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") active_settings = ( settings if isinstance(settings, dict) else facade.load_settings() ) @@ -624,7 +618,7 @@ def recover_aria2_missing_control_file( dependencies: Any = None, ) -> bool: """Restart a stale resolver daemon when a download loses its control file.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") with facade.aria2_lock: other_transfers = [ active_id @@ -670,7 +664,7 @@ def recover_aria2_missing_control_file( def aria2_tell_status(gid: str, *, dependencies: Any = None) -> Dict[str, Any]: """Read an aria2 transfer status with retries for transient resets.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") keys = [ "gid", "status", @@ -711,7 +705,7 @@ def download_file_with_aria2( dependencies: Any = None, ) -> Dict[str, Any]: """Download a file with an aria2c JSON-RPC process.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") settings = facade.load_settings() result = { "success": False, @@ -1025,7 +1019,7 @@ def force_remove_aria2_transfer( dependencies: Any = None, ) -> None: """Request removal of an active aria2 transfer.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") try: facade._aria2_rpc( "aria2.forceRemove", @@ -1041,7 +1035,7 @@ def get_aria2_action_lock( dependencies: Any = None, ) -> Any: """Return the per-download lock used to serialize aria2 actions.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") with facade.aria2_lock: lock = facade.aria2_action_locks.get(download_id) if lock is None: @@ -1058,7 +1052,7 @@ def set_download_progress_status( **updates: Any, ) -> None: """Update a download status while holding the progress lock.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") with facade.download_lock: if download_id in facade.download_progress: facade.download_progress[download_id]["status"] = status @@ -1071,7 +1065,7 @@ def run_aria2_desired_state_worker( dependencies: Any = None, ) -> None: """Apply the latest queued pause/resume request for a transfer.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") while True: with facade.aria2_lock: desired = dict(facade.aria2_desired_states.get(download_id) or {}) @@ -1152,7 +1146,7 @@ def queue_aria2_desired_state( dependencies: Any = None, ) -> Dict[str, Any]: """Queue an aria2 pause/resume state change and start its worker.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") transfer = facade.aria2_transfers.get(download_id) if not transfer or not transfer.get("gid"): return {"success": False, "error": "Download action is not available yet"} @@ -1193,7 +1187,7 @@ def _set_aria2_desired_download_state( *, dependencies: Any = None, ) -> Dict[str, Any]: - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "aria2 backend") if download_id in facade.cancelled_downloads: return {"success": False, "error": "Download is being cancelled"} return facade._queue_aria2_desired_state( diff --git a/core/download/config.py b/core/download/config.py index 57cc14a..f85cf04 100644 --- a/core/download/config.py +++ b/core/download/config.py @@ -3,12 +3,7 @@ import uuid from typing import Any, Dict, Optional - -def _require_dependencies(dependencies: Any) -> Any: - """Return explicitly supplied settings services.""" - if dependencies is None: - raise RuntimeError("download configuration dependencies were not provided") - return dependencies +from .dependencies import require_download_dependencies def generate_download_id() -> str: @@ -22,6 +17,6 @@ def download_backend_from_settings( dependencies: Any = None, ) -> str: """Return the normalized backend selected in the active settings.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "download configuration") active_settings = settings if isinstance(settings, dict) else facade.load_settings() return facade.normalize_download_backend(active_settings.get("download_backend")) diff --git a/core/download/dependencies.py b/core/download/dependencies.py new file mode 100644 index 0000000..a70bc4c --- /dev/null +++ b/core/download/dependencies.py @@ -0,0 +1,10 @@ +"""Shared dependency validation for download components.""" + +from typing import Any + + +def require_download_dependencies(dependencies: Any, component: str) -> Any: + """Return explicitly supplied services for a download component.""" + if dependencies is None: + raise RuntimeError(f"{component} dependencies were not provided") + return dependencies diff --git a/core/download/directories.py b/core/download/directories.py index ea8a904..c4d43f3 100644 --- a/core/download/directories.py +++ b/core/download/directories.py @@ -3,12 +3,7 @@ import os from typing import Any, List, Optional - -def _require_dependencies(dependencies: Any) -> Any: - """Return explicitly supplied services for directory resolution.""" - if dependencies is None: - raise RuntimeError("download directory dependencies were not provided") - return dependencies +from .dependencies import require_download_dependencies def get_download_directory( @@ -18,7 +13,7 @@ def get_download_directory( dependencies: Any = None, ) -> Optional[str]: """Get the appropriate download directory for a model category.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "download directory") folder_paths = facade.folder_paths if folder_paths is None: diff --git a/core/download/huggingface_xet.py b/core/download/huggingface_xet.py index af637a7..0790b2a 100644 --- a/core/download/huggingface_xet.py +++ b/core/download/huggingface_xet.py @@ -5,18 +5,12 @@ from typing import Any, Dict, Optional from urllib.parse import urlparse from ..log_system import create_module_logger +from .dependencies import require_download_dependencies from .state import create_initial_progress log = create_module_logger("core.downloader") -def _require_dependencies(dependencies: Any) -> Any: - """Return explicitly supplied services for Hugging Face Xet.""" - if dependencies is None: - raise RuntimeError("Hugging Face Xet dependencies were not provided") - return dependencies - - class HuggingFaceXetDownloadCancelled(Exception): """Raised by the Xet progress adapter when a download is cancelled.""" @@ -32,7 +26,7 @@ class HuggingFaceXetProgressAdapter: *, dependencies: Any = None, ) -> None: - self.dependencies = _require_dependencies(dependencies) + self.dependencies = require_download_dependencies(dependencies, "Hugging Face Xet") self.download_id = download_id self.total_size = max(0, int(total_size or 0)) self.downloaded = 0 @@ -165,7 +159,7 @@ def run_huggingface_xet_transfer( dependencies: Any = None, ) -> None: """Use detailed native Xet progress when supported, with legacy fallback.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "Hugging Face Xet") import hf_xet from huggingface_hub.file_download import xet_get try: @@ -279,7 +273,7 @@ def download_huggingface_xet( dependencies: Any = None, ) -> Optional[Dict[str, Any]]: """Download Hugging Face Xet files with the official hf_xet transport.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "Hugging Face Xet") validated_url = facade.validate_public_http_url(url) parsed_url = urlparse(validated_url) if not ( diff --git a/core/download/metadata.py b/core/download/metadata.py index f328343..5eef09c 100644 --- a/core/download/metadata.py +++ b/core/download/metadata.py @@ -23,6 +23,7 @@ from ..type_utils import ( first_non_empty, normalize_category_to_model_type, ) +from .dependencies import require_download_dependencies from .validation import _is_sensitive_metadata_key, _strip_sensitive_url_params log = create_module_logger("core.downloader") @@ -32,13 +33,6 @@ _as_list = as_list _first_present = first_non_empty -def _require_dependencies(dependencies: Any) -> Any: - """Return explicitly supplied services for metadata sidecars.""" - if dependencies is None: - raise RuntimeError("metadata sidecar dependencies were not provided") - return dependencies - - def write_model_resolver_metadata( dest_path: str, metadata: Optional[Dict[str, Any]] = None, @@ -49,7 +43,7 @@ def write_model_resolver_metadata( dependencies: Any = None, ) -> Optional[str]: """Write metadata only to the sidecar owned by Model Resolver.""" - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "metadata sidecar") metadata_path = facade.get_model_resolver_sidecar_path(dest_path) try: diff --git a/core/download/previews.py b/core/download/previews.py index 3d9b04c..833e7a3 100644 --- a/core/download/previews.py +++ b/core/download/previews.py @@ -11,6 +11,7 @@ import requests from ..log_system import create_module_logger from ..path_utils import get_filename_from_path from ..type_utils import as_dict, as_list +from .dependencies import require_download_dependencies log = create_module_logger("core.downloader") @@ -35,13 +36,6 @@ MODEL_PREVIEW_EXTENSIONS = ( ) -def _require_dependencies(dependencies: Any) -> Any: - """Return explicitly supplied services for preview downloads.""" - if dependencies is None: - raise RuntimeError("preview download dependencies were not provided") - return dependencies - - def get_existing_model_preview_path(dest_path: str) -> str: """Return the first adjacent preview media file available for a model.""" model_base_path, _model_ext = os.path.splitext(dest_path) @@ -122,7 +116,7 @@ def _download_preview_asset( *, dependencies: Any = None, ) -> bytes: - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "preview download") max_bytes = ( MODEL_PREVIEW_VIDEO_MAX_DOWNLOAD_BYTES if media_type == "video" @@ -200,7 +194,7 @@ def _download_preview_asset_with_system_trust( build_opener, ) - facade = _require_dependencies(dependencies) + facade = require_download_dependencies(dependencies, "preview download") max_bytes = ( MODEL_PREVIEW_VIDEO_MAX_DOWNLOAD_BYTES if media_type == "video" @@ -362,7 +356,7 @@ def create_model_preview( dependencies: Any = None, ) -> str: """Create adjacent preview media using LoRA Manager-compatible behavior.""" - dependencies = _require_dependencies(dependencies) + dependencies = require_download_dependencies(dependencies, "preview download") existing_path = get_existing_model_preview_path(dest_path) source = metadata if isinstance(metadata, dict) else {} preview_url, media_type = _first_model_preview_asset(source) diff --git a/tests/test_download_dependencies.py b/tests/test_download_dependencies.py new file mode 100644 index 0000000..9364ec3 --- /dev/null +++ b/tests/test_download_dependencies.py @@ -0,0 +1,27 @@ +import pytest + +from core.download.dependencies import require_download_dependencies + + +@pytest.mark.parametrize( + ("require_dependencies", "message"), + [ + (require_download_dependencies, "download configuration"), + (require_download_dependencies, "download directory"), + (require_download_dependencies, "metadata sidecar"), + (require_download_dependencies, "Hugging Face Xet"), + (require_download_dependencies, "aria2 backend"), + (require_download_dependencies, "preview download"), + ], +) +def test_download_dependency_guards_preserve_error_and_identity( + require_dependencies, message +): + with pytest.raises( + RuntimeError, + match=rf"{message} dependencies were not provided", + ): + require_dependencies(None, message) + + dependencies = object() + assert require_dependencies(dependencies, message) is dependencies