Coverage for src/photo_download_manager.py: 100%
97 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-05 00:26 +0000
1"""Photo download task management module.
3This module contains utilities for managing photo download tasks
4and parallel execution during photo synchronization.
5"""
7___author___ = "Mandar Patil <mandarons@pm.me>"
9import os
10from concurrent.futures import ThreadPoolExecutor, as_completed
11from threading import Lock
13from src import config_parser, get_logger
14from src.hardlink_registry import HardlinkRegistry
15from src.photo_file_utils import create_hardlink, download_photo_from_server
16from src.photo_path_utils import (
17 _LIVE_VIDEO_SIZES,
18 create_folder_path_if_needed,
19 generate_photo_filename_with_metadata,
20 normalize_file_path,
21 rename_legacy_file_if_exists,
22)
24LOGGER = get_logger()
26# Thread-safe lock for file set operations
27files_lock = Lock()
30class DownloadTaskInfo:
31 """Information about a photo download task."""
33 def __init__(
34 self,
35 photo,
36 file_size: str,
37 photo_path: str,
38 hardlink_source: str | None = None,
39 hardlink_registry: HardlinkRegistry | None = None,
40 ):
41 """Initialize download task info.
43 Args:
44 photo: Photo object from iCloudPy
45 file_size: File size variant (original, medium, thumb, etc.)
46 photo_path: Target path for photo download
47 hardlink_source: Path to existing file for hardlink creation
48 hardlink_registry: Registry for tracking downloaded files
49 """
50 self.photo = photo
51 self.file_size = file_size
52 self.photo_path = photo_path
53 self.hardlink_source = hardlink_source
54 self.hardlink_registry = hardlink_registry
57def get_max_threads_for_download(config) -> int:
58 """Get maximum number of threads for parallel downloads.
60 Args:
61 config: Configuration dictionary
63 Returns:
64 Maximum number of threads to use for downloads
65 """
66 return config_parser.get_app_max_threads(config)
69def generate_photo_path(photo, file_size: str, destination_path: str, folder_format: str | None) -> str:
70 """Generate full file path for photo with legacy file renaming.
72 This function combines path generation, folder creation, and legacy
73 file renaming into a single operation to maintain backward compatibility.
75 Args:
76 photo: Photo object from iCloudPy
77 file_size: File size variant (original, medium, thumb, etc.)
78 destination_path: Base destination path
79 folder_format: strftime format string for folder creation
81 Returns:
82 Normalized full path where photo should be saved
83 """
84 # Generate filename with metadata
85 filename_with_metadata = generate_photo_filename_with_metadata(photo, file_size)
87 # Create folder path if needed
88 final_destination = create_folder_path_if_needed(destination_path, folder_format, photo)
90 # Generate paths for legacy file format handling
91 filename = photo.filename
92 name, extension = filename.rsplit(".", 1) if "." in filename else [filename, ""]
94 # Legacy file paths that need to be renamed
95 file_path = os.path.join(destination_path, filename)
96 file_size_path = os.path.join(
97 destination_path,
98 (f"{'__'.join([name, file_size])}" if extension == "" else f"{'__'.join([name, file_size])}.{extension}"),
99 )
101 # Final path with normalization
102 final_file_path = os.path.join(final_destination, filename_with_metadata)
103 normalized_path = normalize_file_path(final_file_path)
105 # Rename legacy files if they exist
106 rename_legacy_file_if_exists(file_path, normalized_path)
107 rename_legacy_file_if_exists(file_size_path, normalized_path)
109 # Self-heal the earlier .HEIC mislabeling of Live Photo videos: an older
110 # version wrote the paired video with the still's extension. Rename that
111 # file to the corrected path instead of re-downloading it (which would also
112 # leave the broken duplicate behind).
113 if file_size in _LIVE_VIDEO_SIZES and extension:
114 root, _ = os.path.splitext(filename_with_metadata)
115 legacy_mislabeled = normalize_file_path(
116 os.path.join(final_destination, f"{root}.{extension}"),
117 )
118 if legacy_mislabeled != normalized_path:
119 rename_legacy_file_if_exists(legacy_mislabeled, normalized_path)
121 # Handle existing file with different normalization
122 if os.path.isfile(final_file_path) and final_file_path != normalized_path:
123 rename_legacy_file_if_exists(final_file_path, normalized_path)
125 return normalized_path
128def collect_download_task(
129 photo,
130 file_size: str,
131 destination_path: str,
132 files: set[str] | None,
133 folder_format: str | None,
134 hardlink_registry: HardlinkRegistry | None,
135) -> DownloadTaskInfo | None:
136 """Collect photo info for parallel download without immediately downloading.
138 Args:
139 photo: Photo object from iCloudPy
140 file_size: File size variant (original, medium, thumb, etc.)
141 destination_path: Base destination path
142 files: Set to track downloaded files (thread-safe updates)
143 folder_format: strftime format string for folder creation
144 hardlink_registry: Registry for tracking downloaded files
146 Returns:
147 DownloadTaskInfo if photo needs to be processed, None if skipped
148 """
149 # Check if file size exists on server
150 if file_size not in photo.versions:
151 photo_path = generate_photo_path(photo, file_size, destination_path, folder_format)
152 # A missing live_video_* version just means this isn't a Live Photo --
153 # expected for most assets, so log at DEBUG to avoid warning-spam when
154 # live_video_original is in file_sizes. Other sizes warn as before.
155 msg = f"File size {file_size} not found on server. Skipping the photo {photo_path} ..."
156 if file_size.startswith("live_video_"):
157 LOGGER.debug(msg)
158 else:
159 LOGGER.warning(msg)
160 return None
162 # Generate photo path
163 photo_path = generate_photo_path(photo, file_size, destination_path, folder_format)
165 # Thread-safe file set update
166 if files is not None:
167 with files_lock:
168 files.add(photo_path)
170 # Check if photo already exists with correct size
171 from src.photo_file_utils import check_photo_exists
173 if check_photo_exists(photo, file_size, photo_path):
174 return None
176 # Check for existing hardlink source
177 hardlink_source = None
178 if hardlink_registry is not None:
179 hardlink_source = hardlink_registry.get_existing_path(photo.id, file_size)
181 return DownloadTaskInfo(
182 photo=photo,
183 file_size=file_size,
184 photo_path=photo_path,
185 hardlink_source=hardlink_source,
186 hardlink_registry=hardlink_registry,
187 )
190def execute_download_task(task_info: DownloadTaskInfo) -> bool:
191 """Download a single photo or create hardlink as part of parallel execution.
193 Args:
194 task_info: Download task information
196 Returns:
197 True if task completed successfully, False otherwise
198 """
199 LOGGER.debug(f"[Thread] Starting processing of {task_info.photo_path}")
201 try:
202 # Try hardlink first if source exists
203 if task_info.hardlink_source:
204 if create_hardlink(task_info.hardlink_source, task_info.photo_path):
205 LOGGER.debug(f"[Thread] Created hardlink for {task_info.photo_path}")
206 return True
207 else:
208 # Fallback to download if hard link creation fails
209 LOGGER.warning(f"Hard link creation failed, downloading {task_info.photo_path} instead")
211 # Download the photo
212 result = download_photo_from_server(task_info.photo, task_info.file_size, task_info.photo_path)
213 if result and task_info.hardlink_registry is not None:
214 # Register for future hard links if enabled
215 task_info.hardlink_registry.register_photo_path(
216 task_info.photo.id,
217 task_info.file_size,
218 task_info.photo_path,
219 )
220 LOGGER.debug(f"[Thread] Completed download of {task_info.photo_path}")
222 return result
224 except Exception as e:
225 LOGGER.error(f"[Thread] Failed to process {task_info.photo_path}: {e!s}")
226 return False
229def execute_parallel_downloads(download_tasks: list[DownloadTaskInfo], config) -> tuple[int, int]:
230 """Execute download tasks in parallel using thread pool.
232 Args:
233 download_tasks: List of download tasks to execute
234 config: Configuration dictionary for thread settings
236 Returns:
237 Tuple of (successful_downloads, failed_downloads)
238 """
239 if not download_tasks:
240 return 0, 0
242 max_threads = get_max_threads_for_download(config)
244 # Count hardlink tasks vs download tasks for logging
245 hardlink_tasks = sum(1 for task in download_tasks if task.hardlink_source)
246 download_only_tasks = len(download_tasks) - hardlink_tasks
248 if hardlink_tasks > 0:
249 LOGGER.info(
250 f"Starting parallel processing with {max_threads} threads: "
251 f"{hardlink_tasks} hard links, {download_only_tasks} downloads...",
252 )
253 else:
254 LOGGER.info(
255 f"Starting parallel photo downloads with {max_threads} threads for {len(download_tasks)} photos...",
256 )
258 successful_downloads = 0
259 failed_downloads = 0
261 with ThreadPoolExecutor(max_workers=max_threads) as executor:
262 # Submit all download tasks
263 future_to_task = {executor.submit(execute_download_task, task): task for task in download_tasks}
265 # Process completed downloads
266 for future in as_completed(future_to_task):
267 try:
268 result = future.result()
269 if result:
270 successful_downloads += 1
271 else:
272 failed_downloads += 1
273 except Exception as e: # noqa: PERF203
274 LOGGER.error(f"Unexpected error during photo download: {e!s}")
275 failed_downloads += 1
277 LOGGER.info(f"Photo processing complete: {successful_downloads} successful, {failed_downloads} failed")
278 return successful_downloads, failed_downloads