diff -r 5973743a52a1 Doc/library/zipfile.rst --- a/Doc/library/zipfile.rst Sun Mar 06 18:10:58 2011 -0600 +++ b/Doc/library/zipfile.rst Tue Mar 15 03:02:13 2011 +0200 @@ -247,6 +247,16 @@ :meth:`read` on a closed ZipFile will raise a :exc:`RuntimeError`. +.. method:: ZipFile.remove(member) + + Removes the file *member* from the archive. *member* must be the full file + name in the archive or a :class:`ZipInfo` object. The archive must be open in + append mode. Calling :meth:`remove` on a closed ZipFile will raise a + :exc:`RuntimeError`. + + .. versionadded:: 3.3 + + .. method:: ZipFile.testzip() Read all the files in the archive and check their CRC's and file headers. diff -r 5973743a52a1 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() diff -r 5973743a52a1 Lib/zipfile.py --- a/Lib/zipfile.py Sun Mar 06 18:10:58 2011 -0600 +++ b/Lib/zipfile.py Tue Mar 15 03:02:13 2011 +0200 @@ -118,6 +118,9 @@ _FH_FILENAME_LENGTH = 10 _FH_EXTRA_FIELD_LENGTH = 11 +_FHF_HAS_DATA_DESCRIPTOR = 0x8 +dataDescriptorSignature = 0x08074b50 + # The "Zip64 end of central directory locator" structure, magic number, and size structEndArchive64Locator = "<4sLQL" stringEndArchive64Locator = b"PK\x06\x07" @@ -329,7 +332,7 @@ dt = self.date_time dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2] dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2) - if self.flag_bits & 0x08: + if self.flag_bits & _FHF_HAS_DATA_DESCRIPTOR: # Set these to zero because we write them after the file data CRC = compress_size = file_size = 0 else: @@ -658,7 +661,7 @@ class ZipFile: - """ Class with methods to open, read, write, close, list zip files. + """ Class with methods to open, read, write, remove, close, list zip files. z = ZipFile(file, mode="r", compression=ZIP_STORED, allowZip64=False) @@ -963,7 +966,7 @@ # and is used to check the correctness of the password. header = zef_file.read(12) h = list(map(zd, header[0:12])) - if zinfo.flag_bits & 0x8: + if zinfo.flag_bits & _FHF_HAS_DATA_DESCRIPTOR: # compare against the file type from extended local headers check_byte = (zinfo._raw_time >> 8) & 0xff else: @@ -1183,17 +1186,138 @@ self.fp.write(zinfo.FileHeader()) self.fp.write(data) self.fp.flush() - if zinfo.flag_bits & 0x08: + if zinfo.flag_bits & _FHF_HAS_DATA_DESCRIPTOR: # Write CRC and file sizes after the file data self.fp.write(struct.pack(" fileofs: + info.header_offset = info.header_offset - zlen + + # Remove the zipped data + self.fp.seek(fileofs + zlen) + data_after = self.fp.read() + self.fp.seek(fileofs) + self.fp.write(data_after) + new_start_dir = self.start_dir - zlen + self.fp.seek(new_start_dir) + self.fp.truncate() + + # Fix class members with state + self.start_dir = new_start_dir + self._didModify = True + self.filelist.remove(zinfo) + del self.NameToInfo[fname] + def __del__(self): """Call the "close()" method in case the user forgot.""" self.close() + def _central_dir_header(self, zinfo): + dt = zinfo.date_time + dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2] + dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2) + extra = [] + if zinfo.file_size > ZIP64_LIMIT \ + or zinfo.compress_size > ZIP64_LIMIT: + extra.append(zinfo.file_size) + extra.append(zinfo.compress_size) + file_size = 0xffffffff + compress_size = 0xffffffff + else: + file_size = zinfo.file_size + compress_size = zinfo.compress_size + + if zinfo.header_offset > ZIP64_LIMIT: + extra.append(zinfo.header_offset) + header_offset = 0xffffffff + else: + header_offset = zinfo.header_offset + + extra_data = zinfo.extra + if extra: + # Append a ZIP64 field to the extra's + extra_data = struct.pack( + ' ZIP64_LIMIT \ - or zinfo.compress_size > ZIP64_LIMIT: - extra.append(zinfo.file_size) - extra.append(zinfo.compress_size) - file_size = 0xffffffff - compress_size = 0xffffffff - else: - file_size = zinfo.file_size - compress_size = zinfo.compress_size - - if zinfo.header_offset > ZIP64_LIMIT: - extra.append(zinfo.header_offset) - header_offset = 0xffffffff - else: - header_offset = zinfo.header_offset - - extra_data = zinfo.extra - if extra: - # Append a ZIP64 field to the extra's - extra_data = struct.pack( - '