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 Wed Mar 09 22:52:11 2011 +0200 @@ -247,6 +247,14 @@ :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`. + + .. 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 Wed Mar 09 22:52:11 2011 +0200 @@ -1440,11 +1440,83 @@ 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 test_corruption(self): + data = bytes(list(range(100))) + stub_name = "firstfile" + stub_name_2 = "secondfile" + name_to_remove = "foo.txt" + + with zipfile.ZipFile(TESTFN, "w") as zf: + zf.writestr(stub_name, data) + zf.writestr(name_to_remove, "just add a file with a name and some data") + zf.writestr(stub_name_2, data) + + with zipfile.ZipFile(TESTFN, "a") as zf: + zf.remove(name_to_remove) + + with zipfile.ZipFile(TESTFN, "r") as zf: + self.assertEqual(zf.read(stub_name), data) + # this read fails if the pointers weren't corrected + self.assertEqual(zf.read(stub_name_2), data) + + def test_shrinks(self): + fname = "foo.txt" + + with zipfile.ZipFile(TESTFN, "w") as zf: + zf.writestr(fname, "just add a file with a name and some data") + size = os.path.getsize(TESTFN) + with zipfile.ZipFile(TESTFN, "a") as zf: + zf.remove(fname) + + new_size = os.path.getsize(TESTFN) + + # size was 97 bytes vs 139 bytes on my machine btw + self.assertLess(new_size, size) + + 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 Wed Mar 09 22:52:11 2011 +0200 @@ -658,7 +658,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) @@ -1190,6 +1190,46 @@ self.filelist.append(zinfo) self.NameToInfo[zinfo.filename] = zinfo + def remove(self, member): + """Remove a file from the archive. Only works if the ZipFile was opened + with mode 'a'.""" + + if "a" not in self.mode: + raise RuntimeError('remove() requires mode "a"') + if not self.fp: + raise RuntimeError( + "Attempt to modify ZIP archive that was already closed") + + # Make sure we have an info object + if isinstance(member, ZipInfo): + # 'member' is already an info object + zinfo = member + else: + # Get info object for member + zinfo = self.getinfo(member) + + # To remove the member we need its size and location in the archive + fname = zinfo.filename + zlen = len(zinfo.FileHeader()) + zinfo.compress_size + fileofs = zinfo.header_offset + + # Modify all the relevant file pointers + for info in self.infolist(): + if info.header_offset > fileofs: + info.header_offset = info.header_offset - zlen + + # Remove the zipped data + self.fp.seek(fileofs + zlen) + after = self.fp.read() + self.fp.seek(fileofs) + self.fp.write(after) + self.fp.truncate() + + # Fix class members with state + 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()