[276] recursive collection put and get

Author: d-w-moore

irods/collection.py

@@ -43,6 +43,9 @@ def data_objects(self):
             for _, replicas in grouped
         ]
 
+    def get_recursive( self , *arg, **kw):
+        self.manager.get_recursive( self.path, *arg, **kw )
+
     def remove(self, recurse=True, force=False, **options):
         self.manager.remove(self.path, recurse, force, **options)
 

irods/manager/collection_manager.py

@@ -1,4 +1,7 @@
 from __future__ import absolute_import
+import os
+import stat
+import itertools
 from irods.models import Collection
 from irods.manager import Manager
 from irods.message import iRODSMessage, CollectionRequest, FileOpenRequest, ObjCopyRequest, StringStringMap
@@ -9,8 +12,66 @@
 import irods.keywords as kw
 
 
+class GetPathCreationError ( RuntimeError ):
+    """Error denoting the failure to create a new directory for writing.
+    """
+
+def make_writable_dir_if_none_exists( path ):
+    if not os.path.exists(path):
+        os.mkdir(path)
+    if os.path.isdir( path ):
+        os.chmod(path, os.stat(path).st_mode | stat.S_IWUSR)
+    if not os.path.isdir( path ) or not os.access( path, os.W_OK ):
+        raise GetPathCreationError( '{!r} not a writable directory'.format(path) )
+
+try:
+    # Python 2 only
+    from string import maketrans as _maketrans
+except:
+    _maketrans = str.maketrans
+
+_sep2slash = _maketrans(os.path.sep,"/")
+_slash2sep = _maketrans("/",os.path.sep)
+_from_mswin = (lambda path: str.translate(path,_sep2slash)) if os.path.sep != '/' else (lambda x:x)
+_to_mswin =   (lambda path: str.translate(path,_slash2sep)) if os.path.sep != '/' else (lambda x:x)
+
 class CollectionManager(Manager):
 
+    def put_recursive (self, localpath, path, abort_if_not_empty = True, **put_options):
+        c = self.sess.collections.create( path )
+        w = list(itertools.islice(c.walk(), 0, 2))  # dereference first 1 to 2 elements of the walk
+        if abort_if_not_empty and (len(w) > 1 or len(w[0][-1]) > 0):
+            raise RuntimeError('collection {path!r} exists and is non-empty'.format(**locals()))
+        localpath = os.path.normpath(localpath)
+        for my_dir,_,sub_files in os.walk(localpath,topdown=True):
+            dir_without_prefix = os.path.relpath( my_dir, localpath )
+            subcoll = self.sess.collections.create(path if dir_without_prefix == os.path.curdir
+                                                        else path + "/" + _from_mswin(dir_without_prefix))
+            for file_ in sub_files:
+                self.sess.data_objects.put( os.path.join(my_dir,file_), subcoll.path + "/" + file_, **put_options)
+
+
+    def get_recursive (self, path, localpath, abort_if_not_empty = True, **get_options):
+        if os.path.isdir(localpath):
+            w = list(itertools.islice(os.walk(localpath), 0, 2))
+            if abort_if_not_empty and (len(w) > 1 or len(w[0][-1]) > 0):
+                raise RuntimeError('local directory {localpath!r} exists and is non-empty'.format(**locals()))
+        def unprefix (path,prefix=''):
+            return path if not path.startswith(prefix) else path[len(prefix):]
+        c = self.get(path)
+        # TODO ## For a visible percent-complete status:
+        #      #  nbytes = sum(d.size for el in c.walk() for d in el[2])
+        #      ## (Then use eg tqdm module to create progress-bar.)
+        c_prefix = c.path + "/"
+        for coll,_,sub_datas in c.walk(topdown=True):
+            relative_collpath = unprefix (coll.path + "/", c_prefix)
+            new_target_dir = os.path.join(localpath, _to_mswin(relative_collpath))
+            make_writable_dir_if_none_exists( new_target_dir )
+            for data in sub_datas:
+                local_data_path = os.path.join(new_target_dir, data.name)
+                self.sess.data_objects.get( data.path, local_data_path, **get_options )
+
+
     def get(self, path):
         query = self.sess.query(Collection).filter(Collection.name == path)
         try:

irods/test/collection_test.py

@@ -65,6 +65,35 @@ def test_create_recursive_collection(self):
         with self.assertRaises(CollectionDoesNotExist):
             self.sess.collections.get(root_coll_path)
 
+
+    def test_recursive_collection_get_and_put_276(self):
+        try:
+            test_dir = '/tmp/testdir_276'
+            root_coll_path = self.test_coll_path + "/my_deep_collection_276"
+            Depth = 5
+            Objs_per_level = 2
+            Content = 'hello-world'
+            helpers.make_deep_collection(
+                self.sess, root_coll_path, depth=Depth, objects_per_level=Objs_per_level, object_content=Content)
+
+            self.sess.collections.get_recursive(root_coll_path, test_dir)
+            self.sess.collections.put_recursive(test_dir,root_coll_path+"_2")
+            summary = [(elem[0].path,elem[2]) for elem in self.sess.collections.get(root_coll_path+"_2").walk(topdown=True)]
+
+            # check final destination for objects' expected content
+            obj_payloads = [d.open('r').read() for cpath,datas in summary for d in datas]
+            self.assertEqual(obj_payloads, [Content.encode('utf-8')] * (Depth * Objs_per_level))
+
+            # check for increase by 1 in successive collection depths
+            coll_depths =  [elem[0].count('/') for elem in summary]
+            self.assertEqual(list(range(len(coll_depths))), [depth - coll_depths[0] for depth in coll_depths])
+        finally:
+            for c in root_coll_path , root_coll_path + "_2":
+                if self.sess.collections.exists(c):
+                    self.sess.collections.remove(c, force = True)
+            shutil.rmtree(test_dir,ignore_errors = True)
+
+
     def test_remove_deep_collection(self):
         # depth = 100
         depth = 20  # placeholder