remove/delete method for zipfile/tarfile objects

Lib/test/test_zipfile.py

--- a/Lib/test/test_zipfile.py	Sun Mar 06 18:10:58 2011 -0600
+++ b/Lib/test/test_zipfile.py	Tue Mar 15 03:02:13 2011 +0200
@@ -1440,11 +1440,116 @@
         unlink(TESTFN2)
 
 
+class RemoveTests(unittest.TestCase):
+    def test_simple(self):
+        fname = "foo.txt"
+        # remove with fname
+        with zipfile.ZipFile(TESTFN, "w") as zf:
+            zf.writestr(fname, "just add a file with a name and some data")
+            self.assertEqual(zf.infolist()[0].filename, fname)
+        with zipfile.ZipFile(TESTFN, "a") as zf:
+            zf.remove(fname)
+            self.assertEqual(len(zf.infolist()), 0)
+
+        # remove with zipinfo
+        with zipfile.ZipFile(TESTFN, "w") as zf:
+            zf.writestr(fname, "just add a file with a name and some data")
+            self.assertEqual(zf.infolist()[0].filename, fname)
+        with zipfile.ZipFile(TESTFN, "a") as zf:
+            zinfo = zf.getinfo(fname)
+            zf.remove(zinfo)
+            self.assertEqual(len(zf.infolist()), 0)
+
+    def write_three_and_remove_one(self, index):
+        data_list = [bytes([randint(0,255) for i in range(10)]) for i in range(3)]
+        fname_list = ["firstfile", "secondfile", "thirdfile"]
+
+        with zipfile.ZipFile(TESTFN, "w") as zf:
+            for fname, data in zip(fname_list, data_list):
+                zf.writestr(fname, data)
+
+        with zipfile.ZipFile(TESTFN, "a") as zf:
+            zf.remove(fname_list[index])
+
+        i_to_check = list(range(len(data_list)))
+        i_to_check.remove(index)
+        
+        with zipfile.ZipFile(TESTFN, "r") as zf:
+            for i in i_to_check:
+                # the last reads fail if the pointers weren't corrected
+                self.assertEqual(zf.read(fname_list[i]), data_list[i])
+
+    def test_remove_middle(self):
+        self.write_three_and_remove_one(0)
+        self.write_three_and_remove_one(1)
+        self.write_three_and_remove_one(2)
+
+    def test_with_data_descriptor(self):
+        fname = "foo.txt"
+        data = "just add a file with a name and some data"
+        
+        with zipfile.ZipFile(TESTFN, "w") as zf:
+            zinfo = zipfile.ZipInfo(fname)
+            zinfo.flag_bits = zinfo.flag_bits | zipfile._FHF_HAS_DATA_DESCRIPTOR
+            zf.writestr(zinfo, data)
+            
+        with zipfile.ZipFile(TESTFN, "r") as zf:
+            zinfo = zf.getinfo(fname)
+            data_desc_size = zf._get_data_descriptor_size(zinfo)
+            zlen = len(zinfo.FileHeader()) + zinfo.compress_size + len(zf._central_dir_header(zinfo)) + data_desc_size
+
+        size = os.path.getsize(TESTFN)
+        with zipfile.ZipFile(TESTFN, "a") as zf:
+            zf.remove(fname)
+
+        new_size = os.path.getsize(TESTFN)
+
+        size_delta = size - new_size
+        self.assertEqual(size_delta, zlen)
+
+    def test_shrinks(self):
+        fname = "foo.txt"
+        data = "just add a file with a name and some data"
+        
+        with zipfile.ZipFile(TESTFN, "w") as zf:
+            zf.writestr(fname, data)
+            zinfo = zf.getinfo(fname)
+            # ignoring descriptor size because ZipFile by default doesn't use it
+            zlen = len(zinfo.FileHeader()) + zinfo.compress_size + len(zf._central_dir_header(zinfo))
+
+        size = os.path.getsize(TESTFN)
+        with zipfile.ZipFile(TESTFN, "a") as zf:
+            zf.remove(fname)
+
+        new_size = os.path.getsize(TESTFN)
+
+        # size was 22 bytes vs 153 bytes on my machine btw
+        self.assertEqual(new_size, size - zlen)
+
+    def test_verifies_requirements(self):
+        fname = "foo.txt"
+        # test no remove on closed zipfile
+        with self.assertRaises(RuntimeError):
+            zf = zipfile.ZipFile(TESTFN, "w")
+            zf.writestr(fname, "just add a file with a name and some data")
+            zf.close()
+            zf.remove(fname)
+
+        # test no remove without "a"
+        with self.assertRaises(RuntimeError):
+            with zipfile.ZipFile(TESTFN, "w") as zf:
+                zf.writestr(fname, "just add a file with a name and some data")
+                zf.remove(fname)
+
+    def tearDown(self):
+        unlink(TESTFN)
+
+
 def test_main():
     run_unittest(TestsWithSourceFile, TestZip64InSmallFiles, OtherTests,
                  PyZipFileTests, DecryptionTests, TestsWithMultipleOpens,
                  TestWithDirectory, UniversalNewlineTests,
-                 TestsWithRandomBinaryFiles)
+                 TestsWithRandomBinaryFiles, RemoveTests)
 
 if __name__ == "__main__":
     test_main()