diff --git a/frameioclient/client.py b/frameioclient/client.py index d058de14..630ffa84 100644 --- a/frameioclient/client.py +++ b/frameioclient/client.py @@ -322,21 +322,22 @@ def upload(self, asset, file): """ uploader = FrameioUploader(asset, file) uploader.upload() - - def download(self, asset, download_folder): - """ - Download an asset. The method will exit once the file is downloaded. - :Args: - asset (object): The asset object. - download_folder (path): The location to download the file to. + def download(self, asset, download_folder, replace=True): + """ + Download an asset. The method will exit once the file is downloaded. - Example:: + :Args: + asset (object): The asset object. + download_folder (path): The location to download the file to. + replace (bool): Replace the file if it exists. - client.download(asset, "~./Downloads") - """ - downloader = FrameioDownloader(asset, download_folder) - downloader.download() + Example:: + + client.download(asset, "~./Downloads") + """ + downloader = FrameioDownloader(asset, download_folder, replace) + downloader.download() def get_comment(self, comment_id, **kwargs): """ diff --git a/frameioclient/download.py b/frameioclient/download.py index b2ed6cba..8f1bc272 100644 --- a/frameioclient/download.py +++ b/frameioclient/download.py @@ -1,18 +1,62 @@ -import requests -import math import os +import requests +from requests.adapters import HTTPAdapter +from requests.packages.urllib3.util.retry import Retry +from frameioclient.utils import calculate_hash + class FrameioDownloader(object): - def __init__(self, asset, download_folder): + def __init__(self, asset, download_folder, replace): self.asset = asset self.download_folder = download_folder + self.replace = replace + self.attempts = 0 + self.retry_limit = 3 + + self.http_retry_strategy = Retry( + total=3, + backoff_factor=1, + status_forcelist=[408, 500, 502, 503, 504], + method_whitelist=['GET'] + ) def download(self): original_filename = self.asset['name'] final_destination = os.path.join(self.download_folder, original_filename) - + + if os.path.isfile(final_destination) and not self.replace: + try: + raise FileExistsError + except NameError: + raise OSError('File exists') # Python < 3.3 + + adapter = HTTPAdapter(max_retries=self.http_retry_strategy) + http = requests.Session() + http.mount('https://', adapter) + url = self.asset['original'] - r = requests.get(url) - - open(final_destination, 'wb').write(r.content) - \ No newline at end of file + + try: + original_checksum = self.asset['checksums']['xx_hash'] + except (TypeError, KeyError): + original_checksum = None + + while self.attempts < self.retry_limit: + r = http.request('GET', url, stream=True) + + with open(final_destination, 'wb') as handle: + try: + for chunk in r.iter_content(chunk_size=4096): + if chunk: + handle.write(chunk) + except requests.exceptions.ChunkedEncodingError: + self.attempts += 1 + continue + + if not original_checksum: + break + + if calculate_hash(final_destination) == original_checksum: + break + + self.attempts += 1