Coverage for src/photo_download_manager.py: 100%
129 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-08 21:05 +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 DEFAULT_REQUEST_TIMEOUT_SEC, 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 get_default_filename_format,
21 get_file_format,
22 normalize_file_path,
23 rename_legacy_file_if_exists,
24)
26LOGGER = get_logger()
28# Thread-safe lock for file set operations
29files_lock = Lock()
32class DownloadTaskInfo:
33 """Information about a photo download task."""
35 def __init__(
36 self,
37 photo,
38 file_size: str,
39 photo_path: str,
40 hardlink_source: str | None = None,
41 hardlink_registry: HardlinkRegistry | None = None,
42 timeout: int = DEFAULT_REQUEST_TIMEOUT_SEC,
43 ):
44 """Initialize download task info.
46 Args:
47 photo: Photo object from iCloudPy
48 file_size: File size variant (original, medium, thumb, etc.)
49 photo_path: Target path for photo download
50 hardlink_source: Path to existing file for hardlink creation
51 hardlink_registry: Registry for tracking downloaded files
52 timeout: HTTP read timeout for the download, in seconds
53 """
54 self.photo = photo
55 self.file_size = file_size
56 self.photo_path = photo_path
57 self.hardlink_source = hardlink_source
58 self.hardlink_registry = hardlink_registry
59 self.timeout = timeout
62def get_max_threads_for_download(config) -> int:
63 """Get maximum number of threads for parallel downloads.
65 Args:
66 config: Configuration dictionary
68 Returns:
69 Maximum number of threads to use for downloads
70 """
71 return config_parser.get_app_max_threads(config)
74def generate_photo_path(photo, file_size: str, destination_path: str, folder_format: str | None) -> str:
75 """Generate full file path for photo with legacy file renaming.
77 This function combines path generation, folder creation, and legacy
78 file renaming into a single operation to maintain backward compatibility.
80 Args:
81 photo: Photo object from iCloudPy
82 file_size: File size variant (original, medium, thumb, etc.)
83 destination_path: Base destination path
84 folder_format: strftime format string for folder creation
86 Returns:
87 Normalized full path where photo should be saved
88 """
89 # Generate filename with metadata
90 filename_with_metadata = generate_photo_filename_with_metadata(photo, file_size)
92 # Create folder path if needed
93 final_destination = create_folder_path_if_needed(
94 destination_path, folder_format, photo,
95 )
97 # Generate paths for legacy file format handling
98 filename = photo.filename
99 name, extension = filename.rsplit(".", 1) if "." in filename else [filename, ""]
101 # Legacy file paths that need to be renamed
102 file_path = os.path.join(destination_path, filename)
103 file_size_path = os.path.join(
104 destination_path,
105 (f"{'__'.join([name, file_size])}" if extension == "" else f"{'__'.join([name, file_size])}.{extension}"),
106 )
108 # Final path with normalization
109 final_file_path = os.path.join(final_destination, filename_with_metadata)
110 normalized_path = normalize_file_path(final_file_path)
112 # The path the still legitimately occupies under the active naming
113 # scheme. Every legacy rename below moves a file onto *this* variant's
114 # path, and os.rename replaces its target -- so a rename starting from
115 # the still destroys the still and whatever stood at the target. Without
116 # folder_format the destination and the photo's folder are the same
117 # directory, which makes the flat "legacy" paths below identical to the
118 # still's current one.
119 still_path = normalize_file_path(
120 os.path.join(
121 final_destination,
122 generate_photo_filename_with_metadata(photo, "original"),
123 ),
124 )
125 _moves_the_still = file_size in _LIVE_VIDEO_SIZES
127 # Rename legacy files if they exist
128 if not (_moves_the_still and file_path == still_path):
129 rename_legacy_file_if_exists(file_path, normalized_path)
130 if not (_moves_the_still and file_size_path == still_path):
131 rename_legacy_file_if_exists(file_size_path, normalized_path)
133 # Self-heal the earlier .HEIC mislabeling of Live Photo videos: an older
134 # version wrote the paired video with the still's extension. Rename that
135 # file to the corrected path instead of re-downloading it (which would also
136 # leave the broken duplicate behind).
137 if file_size in _LIVE_VIDEO_SIZES and extension:
138 root, _ = os.path.splitext(filename_with_metadata)
139 legacy_mislabeled = normalize_file_path(
140 os.path.join(final_destination, f"{root}.{extension}"),
141 )
142 # Under a naming scheme that keeps the photo's own filename, ``root``
143 # is the still's stem, so the candidate above IS the still's correct
144 # path -- renaming it would move the still onto the video's path and,
145 # because os.rename replaces its target, destroy the video too. The
146 # mislabeled file this heals only exists under a scheme that encodes
147 # the variant in the name, where the two paths cannot coincide.
148 # Confirm the candidate really is the paired video before moving it.
149 # os.rename replaces its target, so an unverified rename destroys two
150 # files at once and logs nothing. CloudKit already reports how long
151 # the video should be, so this costs a stat and no extra I/O: a file
152 # of some other length is not the video, whatever it is named.
153 expected_size = (photo.versions.get(file_size) or {}).get("size")
154 try:
155 candidate_size = (
156 os.path.getsize(legacy_mislabeled)
157 if os.path.isfile(legacy_mislabeled)
158 else None
159 )
160 except OSError: # pragma: no cover - raced away between the two calls
161 candidate_size = None
162 if (
163 legacy_mislabeled not in (normalized_path, still_path)
164 and expected_size is not None
165 and candidate_size == expected_size
166 ):
167 rename_legacy_file_if_exists(legacy_mislabeled, normalized_path)
169 # Handle existing file with different normalization
170 if os.path.isfile(final_file_path) and final_file_path != normalized_path:
171 rename_legacy_file_if_exists(final_file_path, normalized_path)
173 return normalized_path
176def collect_download_task(
177 photo,
178 file_size: str,
179 destination_path: str,
180 files: set[str] | None,
181 folder_format: str | None,
182 hardlink_registry: HardlinkRegistry | None,
183) -> DownloadTaskInfo | None:
184 """Collect photo info for parallel download without immediately downloading.
186 Args:
187 photo: Photo object from iCloudPy
188 file_size: File size variant (original, medium, thumb, etc.)
189 destination_path: Base destination path
190 files: Set to track downloaded files (thread-safe updates)
191 folder_format: strftime format string for folder creation
192 hardlink_registry: Registry for tracking downloaded files
194 Returns:
195 DownloadTaskInfo if photo needs to be processed, None if skipped
196 """
197 # Check if file size exists on server
198 if file_size not in photo.versions:
199 photo_path = generate_photo_path(photo, file_size, destination_path, folder_format)
200 # A missing live_video_* version just means this isn't a Live Photo --
201 # expected for most assets, so log at DEBUG to avoid warning-spam when
202 # live_video_original is in file_sizes. Other sizes warn as before.
203 msg = f"File size {file_size} not found on server. Skipping the photo {photo_path} ..."
204 if file_size.startswith("live_video_"):
205 LOGGER.debug(msg)
206 else:
207 LOGGER.warning(msg)
208 return None
210 # Generate photo path
211 photo_path = generate_photo_path(photo, file_size, destination_path, folder_format)
213 # Check if photo already exists with correct size
214 from src.photo_file_utils import check_photo_exists
216 if check_photo_exists(photo, file_size, photo_path):
217 if files is not None:
218 with files_lock:
219 files.add(photo_path)
220 return None
222 # Filename-collision fallback (``simple`` mode only): a plain
223 # ``IMG_1234.HEIC`` path may already be claimed by a DIFFERENT iCloud
224 # photo that happens to share the human filename. Two collision sources:
225 #
226 # 1. On-disk collision: ``check_photo_exists`` returned False but the
227 # path is occupied -- size mismatch with a photo from a prior sync.
228 # 2. In-flight collision: an earlier call in this same
229 # ``_collect_album_download_tasks`` pass already claimed this plain
230 # path. Without this check the later parallel download silently
231 # overwrites the earlier one and we lose data. ``collect_download_task``
232 # runs sequentially during collection so a plain ``in files`` membership
233 # test under ``files_lock`` is sufficient; we hold the lock only long
234 # enough to read.
235 #
236 # In either case we route this photo to the metadata-suffix filename so
237 # both files coexist and both round-trip stably on future syncs. We use
238 # the ``get_default_filename_format()`` accessor (rather than importing
239 # the module-level constant) so ``set_default_filename_format`` updates
240 # are observed live on every call.
241 # Non-unique naming (simple mode, or a photos.file_format template that may
242 # omit a unique component) can collide; the metadata format cannot.
243 is_non_unique = get_default_filename_format() == "simple" or get_file_format() is not None
244 in_flight_collision = False
245 if is_non_unique and files is not None:
246 with files_lock:
247 in_flight_collision = photo_path in files
248 if is_non_unique and (os.path.isfile(photo_path) or in_flight_collision):
249 # Protect the file we are stepping around. This branch exists to
250 # "preserve both photos", but obsolete-file cleanup deletes anything
251 # absent from ``files`` -- so without this the run writes the suffix
252 # copy and then deletes the very file it just refused to overwrite.
253 # On a library with many repeated filenames that is not a rare edge
254 # case: it deletes a large share of the library every sync, and the
255 # next sync re-downloads it, because freeing the plain path changes
256 # the collision outcome. Only an on-disk collision needs this; an
257 # in-flight one means this run already claimed the path.
258 if files is not None and os.path.isfile(photo_path):
259 with files_lock:
260 files.add(photo_path)
261 suffix_folder = create_folder_path_if_needed(
262 destination_path, folder_format, photo,
263 )
264 suffix_basename = generate_photo_filename_with_metadata(
265 photo, file_size, "metadata",
266 )
267 photo_path = normalize_file_path(os.path.join(suffix_folder, suffix_basename))
268 LOGGER.info(
269 f"Filename collision for {photo.filename} (id={photo.id}); "
270 f"using suffix path {photo_path} to preserve both photos.",
271 )
272 if check_photo_exists(photo, file_size, photo_path):
273 if files is not None:
274 with files_lock:
275 files.add(photo_path)
276 return None
278 if files is not None:
279 with files_lock:
280 files.add(photo_path)
282 # Check for existing hardlink source
283 hardlink_source = None
284 if hardlink_registry is not None:
285 hardlink_source = hardlink_registry.get_existing_path(photo.id, file_size)
287 return DownloadTaskInfo(
288 photo=photo,
289 file_size=file_size,
290 photo_path=photo_path,
291 hardlink_source=hardlink_source,
292 hardlink_registry=hardlink_registry,
293 )
296def execute_download_task(task_info: DownloadTaskInfo) -> bool:
297 """Download a single photo or create hardlink as part of parallel execution.
299 Args:
300 task_info: Download task information
302 Returns:
303 True if task completed successfully, False otherwise
304 """
305 LOGGER.debug(f"[Thread] Starting processing of {task_info.photo_path}")
307 try:
308 # Try hardlink first if source exists
309 if task_info.hardlink_source:
310 if create_hardlink(task_info.hardlink_source, task_info.photo_path):
311 LOGGER.debug(f"[Thread] Created hardlink for {task_info.photo_path}")
312 return True
313 else:
314 # Fallback to download if hard link creation fails
315 LOGGER.warning(
316 f"Hard link creation failed, downloading {task_info.photo_path} instead",
317 )
319 # Download the photo
320 result = download_photo_from_server(
321 task_info.photo,
322 task_info.file_size,
323 task_info.photo_path,
324 timeout=task_info.timeout,
325 )
326 if result and task_info.hardlink_registry is not None:
327 # Register for future hard links if enabled
328 task_info.hardlink_registry.register_photo_path(
329 task_info.photo.id,
330 task_info.file_size,
331 task_info.photo_path,
332 )
333 LOGGER.debug(f"[Thread] Completed download of {task_info.photo_path}")
335 return result
337 except Exception as e:
338 LOGGER.error(f"[Thread] Failed to process {task_info.photo_path}: {e!s}")
339 return False
342def execute_parallel_downloads(
343 download_tasks: list[DownloadTaskInfo], config,
344) -> tuple[int, int]:
345 """Execute download tasks in parallel using thread pool.
347 Args:
348 download_tasks: List of download tasks to execute
349 config: Configuration dictionary for thread settings
351 Returns:
352 Tuple of (successful_downloads, failed_downloads)
353 """
354 if not download_tasks:
355 return 0, 0
357 max_threads = get_max_threads_for_download(config)
359 # Resolved once here rather than per task: without it a stalled CDN
360 # connection blocks its worker thread for the life of the process.
361 timeout = config_parser.get_photos_request_timeout(config)
362 for task in download_tasks:
363 task.timeout = timeout
365 # Count hardlink tasks vs download tasks for logging
366 hardlink_tasks = sum(1 for task in download_tasks if task.hardlink_source)
367 download_only_tasks = len(download_tasks) - hardlink_tasks
369 if hardlink_tasks > 0:
370 LOGGER.info(
371 f"Starting parallel processing with {max_threads} threads: "
372 f"{hardlink_tasks} hard links, {download_only_tasks} downloads...",
373 )
374 else:
375 LOGGER.info(
376 f"Starting parallel photo downloads with {max_threads} threads for {len(download_tasks)} photos...",
377 )
379 successful_downloads = 0
380 failed_downloads = 0
382 with ThreadPoolExecutor(max_workers=max_threads) as executor:
383 # Submit all download tasks
384 future_to_task = {
385 executor.submit(execute_download_task, task): task
386 for task in download_tasks
387 }
389 # Process completed downloads
390 for future in as_completed(future_to_task):
391 try:
392 result = future.result()
393 if result:
394 successful_downloads += 1
395 else:
396 failed_downloads += 1
397 except Exception as e: # noqa: PERF203
398 LOGGER.error(f"Unexpected error during photo download: {e!s}")
399 failed_downloads += 1
401 LOGGER.info(
402 f"Photo processing complete: {successful_downloads} successful, {failed_downloads} failed",
403 )
404 return successful_downloads, failed_downloads