Support recursive payload directory removal

Author: cgaudin

src/bagit/__init__.py

@@ -577,23 +577,31 @@ def add_payload(self, src, dest=None, processes=1):
         self.update_payload(processes=processes)
 
 
-    def remove_payload(self, path, processes=1):
+    def remove_payload(self, path, processes=1, recursive=False):
         """
         Remove a payload file and rebuild manifests.
         """
-        payload_path = os.path.join(self.path, path)
+        payload_path = os.path.normpath(path)
 
-        if not path.startswith("data" + os.sep):
+        if self._path_is_dangerous(payload_path):
+            raise ValueError("Payload path is unsafe: %s" % path)
+
+        if not payload_path.startswith("data" + os.sep):
             raise ValueError("Payload path must start with data/: %s" % path)
 
-        if not os.path.isfile(payload_path):
-            raise ValueError("Payload path must be a file: %s" % path)
+        full_path = os.path.join(self.path, payload_path)
 
-        os.remove(payload_path)
+        if os.path.isfile(full_path):
+            os.remove(full_path)
+        elif os.path.isdir(full_path):
+            if not recursive:
+                raise ValueError("Payload path is a directory: %s" % path)
+            shutil.rmtree(full_path)
+        else:
+            raise ValueError("Payload path does not exist: %s" % path)
 
         self.update_payload(processes=processes)
 
-
     def tagfile_entries(self):
         return dict(
             (key, value)

test.py

@@ -1178,6 +1178,38 @@ def test_add_payload_rejects_unsafe_destination(self):
                 with self.assertRaises(ValueError):
                     bag.add_payload(extra, "../extra.txt")
 
+    def test_remove_payload_directory_recursive(self):
+        bag = bagit.make_bag(self.tmpdir)
+
+        payload_dir = j(self.tmpdir, "data", "extra")
+        os.makedirs(payload_dir)
+
+        with open(j(payload_dir, "one.txt"), "w") as f:
+            f.write("one")
+
+        with open(j(payload_dir, "two.txt"), "w") as f:
+            f.write("two")
+
+        bag.update_payload()
+        self.assertTrue(bag.is_valid())
+
+        bag.remove_payload("data/extra", recursive=True)
+
+        bag = bagit.Bag(self.tmpdir)
+
+        self.assertFalse(os.path.exists(payload_dir))
+        self.assertTrue(bag.is_valid())
+
+    def test_remove_payload_directory_requires_recursive(self):
+        bag = bagit.make_bag(self.tmpdir)
+
+        payload_dir = j(self.tmpdir, "data", "extra")
+        os.makedirs(payload_dir)
+
+        with self.assertRaises(ValueError):
+            bag.remove_payload("data/extra")
+
+
 
 class TestFetch(SelfCleaningTestCase):
     def setUp(self):