Support recursive directory addition to payload

Author: cgaudin

src/bagit/__init__.py

@@ -550,13 +550,13 @@ def add_payload(self, src, dest=None, processes=1):
         """
         Copy a file into the payload directory and rebuild manifests.
         """
-        if not os.path.isfile(src):
-            raise ValueError("Payload source must be a file: %s" % src)
+        if not os.path.isfile(src) and not os.path.isdir(src):
+            raise ValueError("Payload source must be a file or directory: %s" % src)
 
         if dest is None:
             dest = os.path.basename(src)
 
-        payload_dest = os.path.normpath(dest)
+        dest = os.path.normpath(dest)
 
         if os.path.isabs(dest) or dest.startswith("..") or os.path.expanduser(dest) != dest:
             raise ValueError("Payload destination is unsafe: %s" % dest)
@@ -567,12 +567,19 @@ def add_payload(self, src, dest=None, processes=1):
             raise ValueError("Payload destination is unsafe: %s" % dest)
 
         dst = os.path.join(self.path, payload_dest)
-        dst_dir = os.path.dirname(dst)
 
-        if not os.path.isdir(dst_dir):
-            os.makedirs(dst_dir)
+        if os.path.isfile(src):
+            dst_dir = os.path.dirname(dst)
+            if not os.path.isdir(dst_dir):
+                os.makedirs(dst_dir)
 
-        shutil.copy2(src, dst)
+            shutil.copy2(src, dst)
+
+        elif os.path.isdir(src):
+            if os.path.exists(dst):
+                raise ValueError("Payload destination already exists: %s" % dest)
+
+            shutil.copytree(src, dst)
 
         self.update_payload(processes=processes)
 

test.py

@@ -1209,6 +1209,49 @@ def test_remove_payload_directory_requires_recursive(self):
         with self.assertRaises(ValueError):
             bag.remove_payload("data/extra")
 
+    def test_add_payload_directory(self):
+        bag = bagit.make_bag(self.tmpdir)
+
+        src_dir = tempfile.mkdtemp()
+        try:
+            with open(j(src_dir, "one.txt"), "w") as f:
+                f.write("one")
+
+            os.mkdir(j(src_dir, "nested"))
+
+            with open(j(src_dir, "nested", "two.txt"), "w") as f:
+                f.write("two")
+
+            bag.add_payload(src_dir)
+
+            dest_dir = j(self.tmpdir, "data", os.path.basename(src_dir))
+
+            self.assertTrue(os.path.isdir(dest_dir))
+            self.assertTrue(os.path.isfile(j(dest_dir, "one.txt")))
+            self.assertTrue(os.path.isfile(j(dest_dir, "nested", "two.txt")))
+
+            bag = bagit.Bag(self.tmpdir)
+            self.assertTrue(bag.is_valid())
+        finally:
+            shutil.rmtree(src_dir)
+
+    def test_add_payload_directory_with_destination(self):
+        bag = bagit.make_bag(self.tmpdir)
+
+        src_dir = tempfile.mkdtemp()
+        try:
+            with open(j(src_dir, "one.txt"), "w") as f:
+                f.write("one")
+
+            bag.add_payload(src_dir, "imports/sip")
+
+            self.assertTrue(os.path.isdir(j(self.tmpdir, "data", "imports", "sip")))
+            self.assertTrue(os.path.isfile(j(self.tmpdir, "data", "imports", "sip", "one.txt")))
+
+            bag = bagit.Bag(self.tmpdir)
+            self.assertTrue(bag.is_valid())
+        finally:
+            shutil.rmtree(src_dir)
 
 
 class TestFetch(SelfCleaningTestCase):