[#574] Allow for tqdm progress bars to be used

Works for both parallel as well as single threaded gets and puts.

Author: qubixes

irods/manager/data_object_manager.py

@@ -124,7 +124,7 @@ def should_parallelize_transfer( self,
                 open_options[kw.DATA_SIZE_KW] = size
 
 
-    def _download(self, obj, local_path, num_threads, **options):
+    def _download(self, obj, local_path, num_threads, progress_bar, **options):
         """Transfer the contents of a data object to a local file.
 
         Called from get() when a local path is named.
@@ -145,14 +145,17 @@ def _download(self, obj, local_path, num_threads, **options):
                     f.close()
                     if not self.parallel_get( (obj,o), local_file, num_threads = num_threads,
                                               target_resource_name = options.get(kw.RESC_NAME_KW,''),
-                                              data_open_returned_values = data_open_returned_values_):
+                                              data_open_returned_values = data_open_returned_values_,
+                                              progress_bar=progress_bar):
                         raise RuntimeError("parallel get failed")
                 else:
                     for chunk in chunks(o, self.READ_BUFFER_SIZE):
                         f.write(chunk)
+                        if progress_bar is not None:
+                            progress_bar.update(len(chunk))
 
 
-    def get(self, path, local_path = None, num_threads = DEFAULT_NUMBER_OF_THREADS, **options):
+    def get(self, path, local_path = None, num_threads = DEFAULT_NUMBER_OF_THREADS, progress_bar = None, **options):
         """
         Get a reference to the data object at the specified `path'.
 
@@ -163,7 +166,7 @@ def get(self, path, local_path = None, num_threads = DEFAULT_NUMBER_OF_THREADS,
 
         # TODO: optimize
         if local_path:
-            self._download(path, local_path, num_threads = num_threads, **options)
+            self._download(path, local_path, num_threads = num_threads, progress_bar=progress_bar, **options)
 
         query = self.sess.query(DataObject)\
             .filter(DataObject.name == irods_basename(path))\
@@ -180,7 +183,7 @@ def get(self, path, local_path = None, num_threads = DEFAULT_NUMBER_OF_THREADS,
         return iRODSDataObject(self, parent, results)
 
 
-    def put(self, local_path, irods_path, return_data_object = False, num_threads = DEFAULT_NUMBER_OF_THREADS, **options):
+    def put(self, local_path, irods_path, return_data_object = False, num_threads = DEFAULT_NUMBER_OF_THREADS, progress_bar = None, **options):
 
         if self.sess.collections.exists(irods_path):
             obj = iRODSCollection.normalize_path(irods_path, os.path.basename(local_path))
@@ -195,7 +198,7 @@ def put(self, local_path, irods_path, return_data_object = False, num_threads =
                 if not self.parallel_put( local_path, (obj,o), total_bytes = sizelist[0], num_threads = num_threads,
                                           target_resource_name = options.get(kw.RESC_NAME_KW,'') or
                                                                  options.get(kw.DEST_RESC_NAME_KW,''),
-                                          open_options = options ):
+                                          open_options = options, progress_bar = progress_bar):
                     raise RuntimeError("parallel put failed")
             else:
                 with self.open(obj, 'w', **options) as o:
@@ -204,6 +207,8 @@ def put(self, local_path, irods_path, return_data_object = False, num_threads =
                         options[kw.OPR_TYPE_KW] = 1 # PUT_OPR
                     for chunk in chunks(f, self.WRITE_BUFFER_SIZE):
                         o.write(chunk)
+                        if progress_bar is not None:
+                            progress_bar.update(len(chunk))
         if kw.ALL_KW in options:
             repl_options = options.copy()
             repl_options[kw.UPDATE_REPL_KW] = ''
@@ -259,7 +264,8 @@ def parallel_get(self,
                      num_threads = 0,
                      target_resource_name = '',
                      data_open_returned_values = None,
-                     progressQueue = False):
+                     progressQueue = False,
+                     progress_bar = None):
         """Call into the irods.parallel library for multi-1247 GET.
 
         Called from a session.data_objects.get(...) (via the _download method) on
@@ -270,7 +276,8 @@ def parallel_get(self,
         return parallel.io_main( self.sess, data_or_path_, parallel.Oper.GET | (parallel.Oper.NONBLOCKING if async_ else 0), file_,
                                  num_threads = num_threads, target_resource_name = target_resource_name,
                                  data_open_returned_values = data_open_returned_values,
-                                 queueLength = (DEFAULT_QUEUE_DEPTH if progressQueue else 0))
+                                 queueLength = (DEFAULT_QUEUE_DEPTH if progressQueue else 0),
+                                 progress_bar = progress_bar)
 
     def parallel_put(self,
                      file_ ,
@@ -280,6 +287,7 @@ def parallel_put(self,
                      num_threads = 0,
                      target_resource_name = '',
                      open_options = {},
+                     progress_bar = None,
                      progressQueue = False):
         """Call into the irods.parallel library for multi-1247 PUT.
 
@@ -290,6 +298,7 @@ def parallel_put(self,
         return parallel.io_main( self.sess, data_or_path_, parallel.Oper.PUT | (parallel.Oper.NONBLOCKING if async_ else 0), file_,
                                  num_threads = num_threads, total_bytes = total_bytes,  target_resource_name = target_resource_name,
                                  open_options = open_options,
+                                 progress_bar = progress_bar,
                                  queueLength = (DEFAULT_QUEUE_DEPTH if progressQueue else 0)
                                )
 

irods/parallel.py

@@ -223,7 +223,7 @@ def _io_send_bytes_progress (queueObject, item):
 
 COPY_BUF_SIZE = (1024 ** 2) * 4
 
-def _copy_part( src, dst, length, queueObject, debug_info, mgr):
+def _copy_part( src, dst, length, queueObject, debug_info, mgr, progress_bar):
     """
     The work-horse for performing the copy between file and data object.
 
@@ -240,6 +240,8 @@ def _copy_part( src, dst, length, queueObject, debug_info, mgr):
         bytecount += buf_len
         accum += buf_len
         if queueObject and accum and _io_send_bytes_progress(queueObject,accum): accum = 0
+        if progress_bar is not None:
+            progress_bar.update(buf_len)
         if verboseConnection:
             print ("("+debug_info+")",end='',file=sys.stderr)
             sys.stderr.flush()
@@ -301,7 +303,7 @@ def finalize(self):
         self.initial_io.close()
 
 
-def _io_part (objHandle, range_, file_, opr_, mgr_, thread_debug_id = '', queueObject = None ):
+def _io_part (objHandle, range_, file_, opr_, mgr_, thread_debug_id = '', queueObject = None, progress_bar = None):
     """
     Runs in a separate thread to manage the transfer of a range of bytes within the data object.
 
@@ -315,12 +317,12 @@ def _io_part (objHandle, range_, file_, opr_, mgr_, thread_debug_id = '', queueO
     file_.seek(offset)
     if thread_debug_id == '':  # for more succinct thread identifiers while debugging.
         thread_debug_id = str(threading.currentThread().ident)
-    return ( _copy_part (file_, objHandle, length, queueObject, thread_debug_id, mgr_) if Operation.isPut()
-        else _copy_part (objHandle, file_, length, queueObject, thread_debug_id, mgr_) )
+    return ( _copy_part (file_, objHandle, length, queueObject, thread_debug_id, mgr_, progress_bar) if Operation.isPut()
+        else _copy_part (objHandle, file_, length, queueObject, thread_debug_id, mgr_, progress_bar) )
 
 
 def _io_multipart_threaded(operation_ , dataObj_and_IO, replica_token, hier_str, session, fname,
-                           total_size, num_threads, **extra_options):
+                           total_size, num_threads, progress_bar, **extra_options):
     """Called by _io_main.
 
     Carve up (0,total_size) range into `num_threads` parts and initiate a transfer thread for each one.
@@ -366,7 +368,7 @@ def bytes_range_for_thread( i, num_threads, total_bytes,  chunk ):
         mgr.add_io( Io )
         logger.debug(u'target_host = %s', Io.raw.session.pool.account.host)
         if File is None: File = gen_file_handle()
-        futures.append(executor.submit( _io_part, Io, byte_range, File, Operation, mgr, str(counter), queueObject))
+        futures.append(executor.submit( _io_part, Io, byte_range, File, Operation, mgr, str(counter), queueObject, progress_bar))
         counter += 1
         Io = File = None
 
@@ -381,7 +383,7 @@ def bytes_range_for_thread( i, num_threads, total_bytes,  chunk ):
 
 
 
-def io_main( session, Data, opr_, fname, R='', **kwopt):
+def io_main( session, Data, opr_, fname, R='', progress_bar = None, **kwopt):
     """
     The entry point for parallel transfers (multithreaded PUT and GET operations).
 
@@ -395,7 +397,6 @@ def io_main( session, Data, opr_, fname, R='', **kwopt):
     Operation = Oper(opr_)
     d_path = None
     Io = None
-
     if isinstance(Data,tuple):
         (Data, Io) = Data[:2]
 
@@ -468,7 +469,8 @@ def io_main( session, Data, opr_, fname, R='', **kwopt):
     queueLength = kwopt.get('queueLength',0)
     retval = _io_multipart_threaded (Operation, (Data, Io), replica_token, resc_hier, session, fname, total_bytes,
                                      num_threads = num_threads,
-                                     _queueLength = queueLength)
+                                     _queueLength = queueLength,
+                                     progress_bar = progress_bar)
 
     # SessionObject.data_objects.parallel_{put,get} will return:
     #   - immediately with an AsyncNotify instance, if Oper.NONBLOCKING flag is used.