Inmproved download functionality (#79)

Add helper functions to simplify a couple of tasks

Fix FrameioHelpers class

Fix FrameioHelpers import

Fix __init__ imports

Finish fixing helpers...

Don't do multi-part downloads on files smaller than 25 MB

Add download improvements from strombergdev

Fix extra prefix in client.assets.download()

Tweak download function

Fix missing call to _get_path() to set self.destination

Add test to Makefile

Fix download function's keyword via ** keyword args

Don't do multi-part downloads on files smaller than 25 MB

Author: jhodges10

Makefile

@@ -15,3 +15,6 @@ bump-patch:
 
 clean:
 	find . -name "*.pyc" -exec rm -f {} \;
+
+test:
+	cd tests && pipenv run python integration.py

frameioclient/lib/download.py

@@ -8,12 +8,18 @@
 import concurrent.futures
 
 from .utils import Utils
-from .exceptions import DownloadException, WatermarkIDDownloadException, AssetNotFullyUploaded
+from .exceptions import (
+  DownloadException,
+  WatermarkIDDownloadException,
+  AssetNotFullyUploaded,
+  AssetChecksumNotPresent,
+  AssetChecksumMismatch
+)
 
 thread_local = threading.local()
 
 class FrameioDownloader(object):
-  def __init__(self, asset, download_folder, prefix, multi_part=False, concurrency=5, replace=False):
+  def __init__(self, asset, download_folder, prefix=None, replace=False, checksum_verification=True, multi_part=False, concurrency=5):
     self.multi_part = multi_part
     self.asset = asset
     self.asset_type = None
@@ -29,8 +35,10 @@ def __init__(self, asset, download_folder, prefix, multi_part=False, concurrency
     self.prefix = prefix
     self.filename = Utils.normalize_filename(asset["name"])
     self.replace = replace
+    self.checksum_verification = checksum_verification
 
     self._evaluate_asset()
+    self._get_path()
 
   def _evaluate_asset(self):
     if self.asset.get("_type") != "file":
@@ -45,19 +53,39 @@ def _get_session(self):
     return thread_local.session
 
   def _create_file_stub(self):
+    if self.replace == True:
+      os.remove(self.destination) # Remove the file
+      self._create_file_stub() # Create a new stub
+      
     try:
       fp = open(self.destination, "w")
       # fp.write(b"\0" * self.file_size) # Disabled to prevent pre-allocatation of disk space
       fp.close()
-    except FileExistsError as e:
-      if self.replace == True:
-        os.remove(self.destination) # Remove the file
-        self._create_file_stub() # Create a new stub
-      else:
-        print(e)
-        raise e
+
+    except Exception as e:
+      raise e
+
     return True
 
+  def _get_path(self):
+    print("prefix:", self.prefix)
+    if self.prefix != None:
+      self.filename = self.prefix + self.filename
+
+    if self.destination == None:
+      final_destination = os.path.join(self.download_folder, self.filename)
+      self.destination = final_destination
+      
+    return self.destination
+
+  def _get_checksum(self):
+    try:
+      self.original_checksum = self.asset['checksums']['xx_hash']
+    except (TypeError, KeyError):
+      self.original_checksum = None
+    
+    return self.original_checksum
+
   def get_download_key(self):
     try:
       url = self.asset['original']
@@ -84,26 +112,27 @@ def get_download_key(self):
 
     return url
 
-  def get_path(self):
-    if self.prefix != None:
-      self.filename = self.prefix + self.filename
+  def download_handler(self):
+    if os.path.isfile(self.destination) and self.replace != True:
+      try:
+        raise FileExistsError
+      except NameError:
+        raise OSError('File exists')  # Python < 3.3
 
-    if self.destination == None:
-      final_destination = os.path.join(self.download_folder, self.filename)
-      self.destination = final_destination
-      
-    return self.destination
+    url = self.get_download_key()
 
-  def download_handler(self):
-    if os.path.isfile(self.get_path()):
-      print("File already exists at this location.")
-      return self.destination
+    if self.watermarked == True:
+      return self.download(url)
     else:
-      url = self.get_download_key()
-
-      if self.watermarked == True:
+      # Don't use multi-part download for files below 25 MB
+      if self.asset['filesize'] < 26214400:
         return self.download(url)
+      if self.multi_part == True:
+        return self.multi_part_download(url)
       else:
+        # Don't use multi-part download for files below 25 MB
+        if self.asset['filesize'] < 26214400:
+          return self.download(url)
         if self.multi_part == True:
           return self.multi_part_download(url)
         else:
@@ -114,8 +143,17 @@ def download(self, url):
     print("Beginning download -- {} -- {}".format(self.asset["name"], Utils.format_bytes(self.file_size, type="size")))
 
     # Downloading
-    r = requests.get(url)
-    open(self.destination, "wb").write(r.content)
+    session = self._get_session()
+    r = session.get('GET', url, stream=True)
+
+    with open(self.destination, 'wb') as handle:
+      try:
+        # TODO make sure this approach works for SBWM download
+        for chunk in r.iter_content(chunk_size=4096):
+          if chunk:
+            handle.write(chunk)
+      except requests.exceptions.ChunkedEncodingError as e:
+        raise e
 
     download_time = time.time() - start_time
     download_speed = Utils.format_bytes(math.ceil(self.file_size/(download_time)))
@@ -161,7 +199,17 @@ def multi_part_download(self, url):
     download_speed = Utils.format_bytes(math.ceil(self.file_size/(download_time)))
     print("Downloaded {} at {}".format(Utils.format_bytes(self.file_size, type="size"), download_speed))
 
-    return self.destination
+    if self.checksum_verification == True:
+      # Check for checksum, if not present throw error
+      if self._get_checksum() == None:
+        raise AssetChecksumNotPresent
+      else:
+        if Utils.calculate_hash(self.destination) != self.original_checksum:
+          raise AssetChecksumMismatch
+        else:
+          return self.destination
+    else:
+      return self.destination
 
   def download_chunk(self, task):
     # Download a particular chunk

frameioclient/lib/exceptions.py

@@ -40,3 +40,24 @@ def __init__(
     ):
         self.message = message
         super().__init__(self.message)
+
+class AssetChecksumNotPresent(Exception):
+    """Exception raised when there's no checksum present for the Frame.io asset.
+    """
+    def __init__(
+        self, 
+        message="""No checksum found on Frame.io for this asset. This could be because it was uploaded \ 
+            before we introduced the feature, the media pipeline failed to process the asset, or the asset has yet to finish being processed."""
+    ):
+        self.message = message
+        super().__init__(self.message)
+
+class AssetChecksumMismatch(Exception):
+    """Exception raised when the checksum for the downloaded file doesn't match what's found on Frame.io.
+    """
+    def __init__(
+        self, 
+        message="Checksum mismatch, you should re-download the asset to resolve any corrupt bits."
+    ):
+        self.message = message
+        super().__init__(self.message)

frameioclient/service/assets.py

@@ -260,7 +260,7 @@ def upload(self, destination_id, filepath, asset=None):
 
     return asset
 
-  def download(self, asset, download_folder, prefix=None, multi_part=False, concurrency=5, replace=False):
+  def download(self, asset, download_folder, **kwargs):
     """
     Download an asset. The method will exit once the file is downloaded.
 
@@ -272,5 +272,5 @@ def download(self, asset, download_folder, prefix=None, multi_part=False, concur
 
         client.assets.download(asset, "~./Downloads")
     """
-    downloader = FrameioDownloader(asset, download_folder, prefix, multi_part, concurrency)
+    downloader = FrameioDownloader(asset, download_folder, **kwargs)
     return downloader.download_handler()