From abb0873d10e81848bf0171ba12a9cbc5bd65f158 Mon Sep 17 00:00:00 2001
From: goymann <philipp.goymann@mpi-bn.mpg.de>
Date: Fri, 2 Jun 2023 19:34:03 +0000
Subject: [PATCH] add progress bar to downlaod
---
loosolab_s3/cmd.py | 2 +-
loosolab_s3/s3_functions.py | 42 ++++++++++++++++++++++++++++++-------
2 files changed, 36 insertions(+), 8 deletions(-)
diff --git a/loosolab_s3/cmd.py b/loosolab_s3/cmd.py
index dd4c581..3a9089e 100644
--- a/loosolab_s3/cmd.py
+++ b/loosolab_s3/cmd.py
@@ -61,7 +61,7 @@ def run_s3_functions():
else:
utils.check_argparser(args, ['secret', "key", "files", "bucketname"])
- s3.download_s3_objects( args.bucketname, args.files, destination=os.getcwd())
+ s3.download_s3_objects( bucketname, files, destination=os.getcwd(), download_bar=True)
#--------------------------------------------------------------------------------------------------------#
# parse command line arguments:
def argparsefunc():
diff --git a/loosolab_s3/s3_functions.py b/loosolab_s3/s3_functions.py
index b18966f..0fa36f3 100644
--- a/loosolab_s3/s3_functions.py
+++ b/loosolab_s3/s3_functions.py
@@ -85,8 +85,15 @@ class Loosolab_s3:
aws_secret_access_key=credentials["secret"],
endpoint_url=credentials["endpoint"]
)
-
+ client = boto3.client(
+ 's3',
+ config=Config(signature_version=credentials["signature"]),
+ aws_access_key_id=credentials["key"],
+ aws_secret_access_key=credentials["secret"],
+ endpoint_url=credentials["endpoint"]
+ )
self.session = session
+ self.client = client
self.check_s3_credentials()
except Exception as e:
self.__exception_log__("Problem with input!",e)
@@ -208,13 +215,13 @@ class Loosolab_s3:
local_tag = self.etag_checksum(local_file_name)
else:
local_tag = hashlib.md5(open(local_file_name).read().encode('utf-8')).hexdigest()
-
+ print('\n')
self.logger.info('local e-tag of ' + local_file_name +' is : ' + local_tag)
self.logger.info('s3 e-tag of ' + file_name +' is : ' + s3_e_tag)
modBool = local_tag == s3_e_tag
if modBool:
- self.logger.info("Files are not changed!" + str(file_name))
+ self.logger.info("Files are not changed! " + str(file_name))
return modBool
except Exception as e:
@@ -457,7 +464,7 @@ class Loosolab_s3:
return True
#--------------------------------------------------------------------------------------------------------#
- def download_s3_objects(self, bucket_name, file_list, destination=os.getcwd(), compare=True):
+ def download_s3_objects(self, bucket_name, file_list, destination=os.getcwd(), compare=True, download_bar=False):
""" Download files from bucket
Parameter:
----------
@@ -478,11 +485,13 @@ class Loosolab_s3:
modBool = self.compare_s3_etag(bucket_name, file_name, file_path)
if self.check_s3_object_ex(bucket_name, file_name) and not modBool:
if not os.path.exists(os.path.dirname(file_path)):
- print(file_path)
os.makedirs(os.path.dirname(file_path))
self.logger.info('Created directory: ' + os.path.dirname(file_path))
- if self.multipart_upload:
+ if self.multipart_upload and not download_bar:
bucket.download_file(file_name, file_path)
+ elif download_bar:
+
+ self.download(file_path,bucket_name, file_name)
else:
bucket.download_file(file_name, file_path, Config=self.transfer)
try:
@@ -500,4 +509,23 @@ class Loosolab_s3:
else:
bucket.download_file(file_name, file_path, Config=self.transfer)
except Exception as e:
- self.__exception_log__("Could not download file " + file_name , e)
\ No newline at end of file
+ self.__exception_log__("Could not download file " + file_name , e)
+
+#--------------------------------------------------------------------------------------------------------#
+
+ def download(self, local_file_name, s3_bucket, s3_object_key):
+
+ meta_data = self.client.head_object(Bucket=s3_bucket, Key=s3_object_key)
+ total_length = int(meta_data.get('ContentLength', 0))
+ downloaded = 0
+
+ def progress(chunk):
+ nonlocal downloaded
+ downloaded += chunk
+ done = int(50 * downloaded / total_length)
+ sys.stdout.write("\r[%s%s]" % ('=' * done, ' ' * (50-done)) )
+ sys.stdout.flush()
+
+ self.logger.info(f'Downloading {s3_object_key}')
+ with open(local_file_name, 'wb') as f:
+ self.client.download_fileobj(s3_bucket, s3_object_key, f, Callback=progress)
\ No newline at end of file
--
GitLab