Index: Lib/zipfile.py =================================================================== --- Lib/zipfile.py (revision 54190) +++ Lib/zipfile.py (working copy) @@ -568,8 +568,9 @@ def __init__(self, file, mode="r", compression=ZIP_STORED, allowZip64=False): """Open the ZIP file with mode read "r", write "w" or append "a".""" - self._allowZip64 = allowZip64 - self._didModify = False + if mode not in ("r", "w", "a"): + raise RuntimeError, 'ZipFile() requires mode "r", "w", or "a"' + if compression == ZIP_STORED: pass elif compression == ZIP_DEFLATED: @@ -578,6 +579,9 @@ "Compression requires the (missing) zlib module" else: raise RuntimeError, "That compression method is not supported" + + self._allowZip64 = allowZip64 + self._didModify = False self.debug = 0 # Level of printing: 0 through 3 self.NameToInfo = {} # Find file info given name self.filelist = [] # List of ZipInfo instances for archive @@ -720,8 +724,13 @@ def getinfo(self, name): """Return the instance of ZipInfo given 'name'.""" - return self.NameToInfo[name] + info = self.NameToInfo.get(name) + if info is None: + raise RuntimeError, \ + 'There is no item named %r in the archive' % name + return info + def setpassword(self, pwd): """Set default password for encrypted files.""" self.pwd = pwd @@ -824,6 +833,10 @@ def write(self, filename, arcname=None, compress_type=None): """Put the bytes from filename into the archive under the name arcname.""" + if not self.fp: + raise RuntimeError, \ + "Attempt to write to ZIP archive that was already closed" + st = os.stat(filename) mtime = time.localtime(st.st_mtime) date_time = mtime[0:6] @@ -896,6 +909,11 @@ zinfo.compress_type = self.compression else: zinfo = zinfo_or_arcname + + if not self.fp: + raise RuntimeError, \ + "Attempt to write to ZIP archive that was already closed" + zinfo.file_size = len(bytes) # Uncompressed size zinfo.header_offset = self.fp.tell() # Start of header bytes self._writecheck(zinfo) Index: Lib/test/test_zipfile.py =================================================================== --- Lib/test/test_zipfile.py (revision 54190) +++ Lib/test/test_zipfile.py (working copy) @@ -13,12 +13,12 @@ from test.test_support import TESTFN, run_unittest TESTFN2 = TESTFN + "2" -FIXEDTEST_SIZE = 10 +FIXEDTEST_SIZE = 1000 class TestsWithSourceFile(unittest.TestCase): def setUp(self): - self.line_gen = ("Zipfile test line %d. random float: %f" % (i, random()) - for i in xrange(FIXEDTEST_SIZE)) + self.line_gen = ["Zipfile test line %d. random float: %f" % (i, random()) + for i in xrange(FIXEDTEST_SIZE)] self.data = '\n'.join(self.line_gen) + '\n' # Make a source file with some lines @@ -238,6 +238,63 @@ self.assertEqual(zipfp.namelist(), ["absolute"]) zipfp.close() + def testAppendToZipFile(self): + # Test appending to an existing zipfile + zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED) + zipfp.write(TESTFN, TESTFN) + zipfp.close() + zipfp = zipfile.ZipFile(TESTFN2, "a", zipfile.ZIP_STORED) + zipfp.writestr("strfile", self.data) + self.assertEqual(zipfp.namelist(), [TESTFN, "strfile"]) + zipfp.close() + + def testAppendToNonZipFile(self): + # Test appending to an existing file that is not a zipfile + # NOTE: this test fails if len(d) < 22 because of the first + # line "fpin.seek(-22, 2)" in _EndRecData + d = 'I am not a ZipFile!'*10 + f = file(TESTFN2, 'wb') + f.write(d) + f.close() + zipfp = zipfile.ZipFile(TESTFN2, "a", zipfile.ZIP_STORED) + zipfp.write(TESTFN, TESTFN) + zipfp.close() + + f = file(TESTFN2, 'rb') + f.seek(len(d)) + zipfp = zipfile.ZipFile(f, "r") + self.assertEqual(zipfp.namelist(), [TESTFN]) + zipfp.close() + f.close() + + def test_WriteDefaultName(self): + # Check that calling ZipFile.write without arcname specified produces the expected result + zipfp = zipfile.ZipFile(TESTFN2, "w") + zipfp.write(TESTFN) + self.assertEqual(zipfp.read(TESTFN), file(TESTFN).read()) + zipfp.close() + + def test_PerFileCompression(self): + # Check that files within a Zip archive can have different compression options + zipfp = zipfile.ZipFile(TESTFN2, "w") + zipfp.write(TESTFN, 'storeme', zipfile.ZIP_STORED) + zipfp.write(TESTFN, 'deflateme', zipfile.ZIP_DEFLATED) + sinfo = zipfp.getinfo('storeme') + dinfo = zipfp.getinfo('deflateme') + self.assertEqual(sinfo.compress_type, zipfile.ZIP_STORED) + self.assertEqual(dinfo.compress_type, zipfile.ZIP_DEFLATED) + zipfp.close() + + def test_WriteToReadonly(self): + # Check that trying to call write() on a readonly ZipFile object + # raises a RuntimeError + zipf = zipfile.ZipFile(TESTFN2, mode="w") + zipf.writestr("somefile.txt", "bogus") + zipf.close() + zipf = zipfile.ZipFile(TESTFN2, mode="r") + self.assertRaises(RuntimeError, zipf.write, TESTFN) + zipf.close() + def tearDown(self): os.remove(TESTFN) os.remove(TESTFN2) @@ -360,7 +417,6 @@ self.assertEqual(zipfp.namelist(), ["absolute"]) zipfp.close() - def tearDown(self): zipfile.ZIP64_LIMIT = self._limit os.remove(TESTFN) @@ -431,6 +487,11 @@ finally: shutil.rmtree(TESTFN2) + def testWriteNonPyfile(self): + zipfp = zipfile.PyZipFile(TemporaryFile(), "w") + file(TESTFN, 'w').write('most definitely not a python file') + self.assertRaises(RuntimeError, zipfp.writepy, TESTFN) + os.remove(TESTFN) class OtherTests(unittest.TestCase): @@ -516,8 +577,62 @@ # a RuntimeError, and so should calling .testzip. An earlier # version of .testzip would swallow this exception (and any other) # and report that the first file in the archive was corrupt. + self.assertRaises(RuntimeError, zipf.read, "foo.txt") + self.assertRaises(RuntimeError, zipf.open, "foo.txt") self.assertRaises(RuntimeError, zipf.testzip) + self.assertRaises(RuntimeError, zipf.writestr, "bogus.txt", "bogus") + file(TESTFN, 'w').write('zipfile test data') + self.assertRaises(RuntimeError, zipf.write, TESTFN) + os.remove(TESTFN) + def test_BadConstructorMode(self): + # Check that bad modes passed to ZipFile constructor are caught + self.assertRaises(RuntimeError, zipfile.ZipFile, TESTFN, "q") + + def test_BadOpenMode(self): + # Check that bad modes passed to ZipFile.open are caught + zipf = zipfile.ZipFile(TESTFN, mode="w") + zipf.writestr("foo.txt", "O, for a Muse of Fire!") + zipf.close() + zipf = zipfile.ZipFile(TESTFN, mode="r") + # read the data to make sure the file is there + zipf.read("foo.txt") + self.assertRaises(RuntimeError, zipf.open, "foo.txt", "q") + zipf.close() + os.unlink(TESTFN) + + def test_Read0(self): + # Check that calling read(0) on a ZipExtFile object returns an empty + # string and doesn't advance file pointer + zipf = zipfile.ZipFile(TESTFN, mode="w") + zipf.writestr("foo.txt", "O, for a Muse of Fire!") + # read the data to make sure the file is there + f = zipf.open("foo.txt") + for i in xrange(FIXEDTEST_SIZE): + self.assertEqual(f.read(0), '') + + self.assertEqual(f.read(), "O, for a Muse of Fire!") + zipf.close() + os.unlink(TESTFN) + + def test_OpenNonexistentItem(self): + # Check that attempting to call open() for an item that doesn't + # exist in the archive raises a RuntimeError + zipf = zipfile.ZipFile(TESTFN, mode="w") + self.assertRaises(RuntimeError, zipf.open, "foo.txt", "r") + os.unlink(TESTFN) + + def test_BadCompressionMode(self): + # Check that bad compression methods passed to ZipFile.open are caught + self.assertRaises(RuntimeError, zipfile.ZipFile, TESTFN, "w", -1) + + def test_NullByteInFilename(self): + # Check that a filename containing a null byte is properly terminated + zipf = zipfile.ZipFile(TESTFN, mode="w") + zipf.writestr("foo.txt\x00qqq", "O, for a Muse of Fire!") + self.assertEqual(zipf.namelist(), ['foo.txt']) + os.unlink(TESTFN) + class DecryptionTests(unittest.TestCase): # This test checks that ZIP decryption works. Since the library does not # support encryption at the moment, we use a pre-generated encrypted