Fix KeyError #859 (#861)

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- fix: KeyError alternative [#859](https://github.com/icloud-photos-downloader/icloud_photos_downloader/issues/859)
+
 ## 1.19.0 (2024-05-31)
 
 - fix: release notes [#849](https://github.com/icloud-photos-downloader/icloud_photos_downloader/issues/849)

src/icloudpd/autodelete.py

@@ -9,6 +9,7 @@
 from icloudpd.paths import local_download_path
 from pyicloud_ipd.services.photos import PhotoLibrary
 from pyicloud_ipd.utils import disambiguate_filenames
+from pyicloud_ipd.version_size import AssetVersionSize, VersionSize
 
 
 def delete_file(logger: logging.Logger, path: str) -> bool:
@@ -30,7 +31,7 @@ def autodelete_photos(
         library_object: PhotoLibrary,
         folder_structure: str,
         directory: str,
-        _sizes: Sequence[str]) -> None:
+        _sizes: Sequence[AssetVersionSize]) -> None:
     """
     Scans the "Recently Deleted" folder and deletes any matching files
     from the download directory.
@@ -67,13 +68,14 @@ def autodelete_photos(
         download_dir = os.path.join(directory, date_path)
 
         paths:Set[str] = set({})
+        _size:VersionSize
         for _size, _version in disambiguate_filenames(media.versions, _sizes).items():
-            if _size in ["alternative", "adjusted"]:
+            if _size in [AssetVersionSize.ALTERNATIVE, AssetVersionSize.ADJUSTED]:
                 paths.add(os.path.normpath(
                     local_download_path(
                         _version["filename"], download_dir)))
         for _size, _version in media.versions.items():
-            if _size not in ["alternative", "adjusted"]:
+            if _size not in [AssetVersionSize.ALTERNATIVE, AssetVersionSize.ADJUSTED]:
                 paths.add(os.path.normpath(
                     local_download_path(
                         _version["filename"], download_dir)))

src/icloudpd/base.py

@@ -32,6 +32,7 @@
 from multiprocessing import freeze_support
 
 from pyicloud_ipd.utils import compose, disambiguate_filenames, identity
+from pyicloud_ipd.version_size import AssetVersionSize, LivePhotoVersionSize
 freeze_support()  # fixing tqdm on macos
 
 def build_filename_cleaner(_ctx: click.Context, _param: click.Parameter, is_keep_unicode: bool) -> Callable[[str], str]:
@@ -71,6 +72,32 @@ def raw_policy_generator(_ctx: click.Context, _param: click.Parameter, raw_polic
     else:
         raise ValueError(f"policy was provided with unsupported value of '{raw_policy}'")
 
+def size_generator(_ctx: click.Context, _param: click.Parameter, sizes: Sequence[str]) -> Sequence[AssetVersionSize]:
+    def _map(size: str) -> AssetVersionSize:
+        if size == "original":
+            return AssetVersionSize.ORIGINAL
+        elif size == "adjusted":
+            return AssetVersionSize.ADJUSTED
+        elif size == "alternative":
+            return AssetVersionSize.ALTERNATIVE
+        elif size == "medium":
+            return AssetVersionSize.MEDIUM
+        elif size == "thumb":
+            return AssetVersionSize.THUMB
+        else:
+            raise ValueError(f"size was provided with unsupported value of '{size}'")    
+    return [_map(_s) for _s in sizes]
+
+def lp_size_generator(_ctx: click.Context, _param: click.Parameter, size: str) -> LivePhotoVersionSize:
+    if size == "original":
+        return LivePhotoVersionSize.ORIGINAL
+    elif size == "medium":
+        return LivePhotoVersionSize.MEDIUM
+    elif size == "thumb":
+        return LivePhotoVersionSize.THUMB
+    else:
+        raise ValueError(f"size was provided with unsupported value of '{size}'")    
+
 # Must import the constants object so that we can mock values in tests.
 
 CONTEXT_SETTINGS = {"help_option_names": ["-h", "--help"]}
@@ -114,13 +141,15 @@ def raw_policy_generator(_ctx: click.Context, _param: click.Parameter, raw_polic
     default=["original"],
     multiple=True,
     show_default=True,
+    callback=size_generator,
 )
 @click.option(
     "--live-photo-size",
     help="Live Photo video size to download",
     type=click.Choice(["original", "medium", "thumb"]),
     default="original",
     show_default=True,
+    callback=lp_size_generator,
 )
 @click.option(
     "--recent",
@@ -315,8 +344,8 @@ def main(
         password: Optional[str],
         auth_only: bool,
         cookie_directory: str,
-        size: Sequence[str],
-        live_photo_size: str,
+        size: Sequence[AssetVersionSize],
+        live_photo_size: LivePhotoVersionSize,
         recent: Optional[int],
         until_found: Optional[int],
         album: str,
@@ -451,12 +480,12 @@ def download_builder(
         skip_videos: bool,
         folder_structure: str,
         directory: str,
-        size: Sequence[str],
+        size: Sequence[AssetVersionSize],
         force_size: bool,
         only_print_filenames: bool,
         set_exif_datetime: bool,
         skip_live_photos: bool,
-        live_photo_size: str,
+        live_photo_size: LivePhotoVersionSize,
         dry_run: bool
         ) -> Callable[[PyiCloudService], Callable[[Counter, PhotoAsset], bool]]:
     """factory for downloader"""
@@ -532,17 +561,17 @@ def download_photo_(counter: Counter, photo: PhotoAsset) -> bool:
             success = False
 
             for download_size in size:
-                if download_size not in versions and download_size != "original":
+                if download_size not in versions and download_size != AssetVersionSize.ORIGINAL:
                     if force_size:
                         logger.error(
                             "%s size does not exist for %s. Skipping...",
-                            download_size,
+                            download_size.value,
                             photo.filename
                         )
                         return False
-                    if "original" in size:
+                    if AssetVersionSize.ORIGINAL in size:
                         continue    # that should avoid double download for original
-                    download_size = "original"
+                    download_size = AssetVersionSize.ORIGINAL
 
                 version = versions[download_size]
                 filename = version["filename"]
@@ -552,11 +581,11 @@ def download_photo_(counter: Counter, photo: PhotoAsset) -> bool:
 
                 original_download_path = None
                 file_exists = os.path.isfile(download_path)
-                if not file_exists and download_size == "original":
+                if not file_exists and download_size == AssetVersionSize.ORIGINAL:
                     # Deprecation - We used to download files like IMG_1234-original.jpg,
                     # so we need to check for these.
                     # Now we match the behavior of iCloud for Windows: IMG_1234.jpg
-                    original_download_path = (f"-{download_size}.").join(
+                    original_download_path = (f"-original.").join(
                         download_path.rsplit(".", 1)
                     )
                     file_exists = os.path.isfile(original_download_path)
@@ -623,7 +652,7 @@ def download_photo_(counter: Counter, photo: PhotoAsset) -> bool:
 
             # Also download the live photo if present
             if not skip_live_photos:
-                lp_size = live_photo_size + "Video"
+                lp_size = live_photo_size
                 if lp_size in photo.versions:
                     version = photo.versions[lp_size]
                     lp_filename = version["filename"]
@@ -799,7 +828,7 @@ def core(
         password: Optional[str],
         auth_only: bool,
         cookie_directory: str,
-        size: Sequence[str],
+        size: Sequence[AssetVersionSize],
         recent: Optional[int],
         until_found: Optional[int],
         album: str,
@@ -966,7 +995,7 @@ def core(
                 ("Downloading %s %s" +
                  " photo%s%s to %s ..."),
                 photos_count_str,
-                ",".join(size),
+                ",".join([_s.value for _s in size]),
                 plural_suffix,
                 video_suffix,
                 directory

src/icloudpd/download.py

@@ -15,6 +15,7 @@
 
 # Import the constants object so that we can mock WAIT_SECONDS in tests
 from icloudpd import constants
+from pyicloud_ipd.version_size import VersionSize
 
 
 def update_mtime(created: datetime.datetime, download_path: str) -> None:
@@ -105,7 +106,7 @@ def download_media(
         photo: PhotoAsset,
         download_path: str,
         version: Dict[str, Any],
-        size: str) -> bool:
+        size: VersionSize) -> bool:
     """Download the photo to path, with retries and error handling"""
 
     mkdirs_local = mkdirs_for_path_dry_run if dry_run else mkdirs_for_path
@@ -124,7 +125,7 @@ def download_media(
             logger.error(
                 "Could not find URL to download %s for size %s",
                 version["filename"],
-                size
+                size.value
             )
             break
 

src/pyicloud_ipd/services/photos.py

@@ -6,7 +6,7 @@
 import re
 
 from datetime import datetime
-from typing import Any, Callable, Dict, Generator, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, Generator, Optional, Sequence, Tuple, Union
 import typing
 
 from requests import Response
@@ -19,6 +19,7 @@
 
 from pyicloud_ipd.raw_policy import RawTreatmentPolicy
 from pyicloud_ipd.session import PyiCloudSession
+from pyicloud_ipd.version_size import AssetVersionSize, LivePhotoVersionSize, VersionSize
 
 logger = logging.getLogger(__name__)
 
@@ -526,7 +527,7 @@ def __init__(self, service:PhotosService, master_record: Dict[str, Any], asset_r
         self._master_record = master_record
         self._asset_record = asset_record
 
-        self._versions: Optional[Dict[str, Dict[str, Any]]] = None
+        self._versions: Optional[Dict[VersionSize, Dict[str, Any]]] = None
 
     ITEM_TYPES = {
         u"public.heic": u"image",
@@ -566,21 +567,28 @@ def __init__(self, service:PhotosService, master_record: Dict[str, Any], asset_r
         u'com.olympus.or-raw-image': u"ORF",
     }
 
-    PHOTO_VERSION_LOOKUP = {
-        u"original": u"resOriginal",
-        u"alternative": u"resOriginalAlt",
-        u"medium": u"resJPEGMed",
-        u"thumb": u"resJPEGThumb",
-        u"adjusted": u"resJPEGFull",
-        u"originalVideo": u"resOriginalVidCompl",
-        u"mediumVideo": u"resVidMed",
-        u"thumbVideo": u"resVidSmall",
+    PHOTO_VERSION_LOOKUP: Dict[VersionSize, str] = {
+        AssetVersionSize.ORIGINAL: u"resOriginal",
+        AssetVersionSize.ALTERNATIVE: u"resOriginalAlt",
+        AssetVersionSize.MEDIUM: u"resJPEGMed",
+        AssetVersionSize.THUMB: u"resJPEGThumb",
+        AssetVersionSize.ADJUSTED: u"resJPEGFull",
+        LivePhotoVersionSize.ORIGINAL: u"resOriginalVidCompl",
+        LivePhotoVersionSize.MEDIUM: u"resVidMed",
+        LivePhotoVersionSize.THUMB: u"resVidSmall",
     }
 
-    VIDEO_VERSION_LOOKUP = {
-        u"original": u"resOriginal",
-        u"medium": u"resVidMed",
-        u"thumb": u"resVidSmall"
+    VIDEO_VERSION_LOOKUP: Dict[VersionSize, str] = {
+        AssetVersionSize.ORIGINAL: u"resOriginal",
+        AssetVersionSize.MEDIUM: u"resVidMed",
+        AssetVersionSize.THUMB: u"resVidSmall"
+    }
+
+    VERSION_FILENAME_SUFFIX_LOOKUP: Dict[VersionSize, str] = {
+        AssetVersionSize.MEDIUM: u"medium",
+        AssetVersionSize.THUMB: u"thumb",
+        LivePhotoVersionSize.MEDIUM: u"medium",
+        LivePhotoVersionSize.THUMB: u"thumb",
     }
 
     @property
@@ -654,11 +662,11 @@ def item_type_extension(self) -> str:
         return 'unknown'
 
     @property
-    def versions(self) -> Dict[str, Dict[str, Any]]:
+    def versions(self) -> Dict[VersionSize, Dict[str, Any]]:
         if not self._versions:
-            _versions: Dict[str, Dict[str, Any]] = {}
+            _versions: Dict[VersionSize, Dict[str, Any]] = {}
             if self.item_type == "movie":
-                typed_version_lookup = self.VIDEO_VERSION_LOOKUP
+                typed_version_lookup: Dict[VersionSize, str] = self.VIDEO_VERSION_LOOKUP
             else:
                 typed_version_lookup = self.PHOTO_VERSION_LOOKUP
 
@@ -708,23 +716,20 @@ def versions(self) -> Dict[str, Dict[str, Any]]:
                         _f, _e = os.path.splitext(version["filename"])
                         version["filename"] = _f + "." + self.ITEM_TYPE_EXTENSIONS.get(version["type"], _e[1:])
 
-                    # add size
-                    if "Video" in key:
-                        _size_cleaned = key[:-5]
-                    else:
-                        _size_cleaned = key
-                    if _size_cleaned not in ["original", "adjusted", "alternative"]:
+                    # add size suffix
+                    if key in self.VERSION_FILENAME_SUFFIX_LOOKUP:
+                        _size_suffix = self.VERSION_FILENAME_SUFFIX_LOOKUP[key]
                         _f, _e = os.path.splitext(version["filename"])
-                        version["filename"] = _f + f"-{_size_cleaned}" + _e
+                        version["filename"] = _f + f"-{_size_suffix}" + _e
 
                     _versions[key] = version
 
             # swap original & alternative according to swap_raw_policy
-            if "alternative" in _versions and (("raw" in _versions["alternative"]["type"] and self._service.raw_policy == RawTreatmentPolicy.AS_ORIGINAL) or ("raw" in _versions["original"]["type"] and self._service.raw_policy == RawTreatmentPolicy.AS_ALTERNATIVE)):
-                _a = dict(_versions["alternative"])
-                _o = dict(_versions["original"])
-                _versions["alternative"] = _o
-                _versions["original"] = _a
+            if AssetVersionSize.ALTERNATIVE in _versions and (("raw" in _versions[AssetVersionSize.ALTERNATIVE]["type"] and self._service.raw_policy == RawTreatmentPolicy.AS_ORIGINAL) or ("raw" in _versions[AssetVersionSize.ORIGINAL]["type"] and self._service.raw_policy == RawTreatmentPolicy.AS_ALTERNATIVE)):
+                _a = dict(_versions[AssetVersionSize.ALTERNATIVE])
+                _o = dict(_versions[AssetVersionSize.ORIGINAL])
+                _versions[AssetVersionSize.ALTERNATIVE] = _o
+                _versions[AssetVersionSize.ORIGINAL] = _a
 
             self._versions = _versions