test_zipfile.py revision 6d6b53cab8099c07aab70887aeef238b0c8eb94f
1# We can test part of the module without zlib.
2try:
3    import zlib
4except ImportError:
5    zlib = None
6
7import os
8import sys
9import time
10import shutil
11import struct
12import zipfile
13import unittest
14
15from StringIO import StringIO
16from tempfile import TemporaryFile
17from random import randint, random
18from unittest import skipUnless
19
20from test.test_support import TESTFN, run_unittest, findfile, unlink
21
22TESTFN2 = TESTFN + "2"
23TESTFNDIR = TESTFN + "d"
24FIXEDTEST_SIZE = 1000
25
26SMALL_TEST_DATA = [('_ziptest1', '1q2w3e4r5t'),
27                   ('ziptest2dir/_ziptest2', 'qawsedrftg'),
28                   ('/ziptest2dir/ziptest3dir/_ziptest3', 'azsxdcfvgb'),
29                   ('ziptest2dir/ziptest3dir/ziptest4dir/_ziptest3', '6y7u8i9o0p')]
30
31
32class TestsWithSourceFile(unittest.TestCase):
33    def setUp(self):
34        self.line_gen = ["Zipfile test line %d. random float: %f" % (i, random())
35                         for i in xrange(FIXEDTEST_SIZE)]
36        self.data = '\n'.join(self.line_gen) + '\n'
37
38        # Make a source file with some lines
39        with open(TESTFN, "wb") as fp:
40            fp.write(self.data)
41
42    def make_test_archive(self, f, compression):
43        # Create the ZIP archive
44        with zipfile.ZipFile(f, "w", compression) as zipfp:
45            zipfp.write(TESTFN, "another.name")
46            zipfp.write(TESTFN, TESTFN)
47            zipfp.writestr("strfile", self.data)
48
49    def zip_test(self, f, compression):
50        self.make_test_archive(f, compression)
51
52        # Read the ZIP archive
53        with zipfile.ZipFile(f, "r", compression) as zipfp:
54            self.assertEqual(zipfp.read(TESTFN), self.data)
55            self.assertEqual(zipfp.read("another.name"), self.data)
56            self.assertEqual(zipfp.read("strfile"), self.data)
57
58            # Print the ZIP directory
59            fp = StringIO()
60            stdout = sys.stdout
61            try:
62                sys.stdout = fp
63                zipfp.printdir()
64            finally:
65                sys.stdout = stdout
66
67            directory = fp.getvalue()
68            lines = directory.splitlines()
69            self.assertEqual(len(lines), 4) # Number of files + header
70
71            self.assertTrue('File Name' in lines[0])
72            self.assertTrue('Modified' in lines[0])
73            self.assertTrue('Size' in lines[0])
74
75            fn, date, time_, size = lines[1].split()
76            self.assertEqual(fn, 'another.name')
77            self.assertTrue(time.strptime(date, '%Y-%m-%d'))
78            self.assertTrue(time.strptime(time_, '%H:%M:%S'))
79            self.assertEqual(size, str(len(self.data)))
80
81            # Check the namelist
82            names = zipfp.namelist()
83            self.assertEqual(len(names), 3)
84            self.assertTrue(TESTFN in names)
85            self.assertTrue("another.name" in names)
86            self.assertTrue("strfile" in names)
87
88            # Check infolist
89            infos = zipfp.infolist()
90            names = [i.filename for i in infos]
91            self.assertEqual(len(names), 3)
92            self.assertTrue(TESTFN in names)
93            self.assertTrue("another.name" in names)
94            self.assertTrue("strfile" in names)
95            for i in infos:
96                self.assertEqual(i.file_size, len(self.data))
97
98            # check getinfo
99            for nm in (TESTFN, "another.name", "strfile"):
100                info = zipfp.getinfo(nm)
101                self.assertEqual(info.filename, nm)
102                self.assertEqual(info.file_size, len(self.data))
103
104            # Check that testzip doesn't raise an exception
105            zipfp.testzip()
106
107    def test_stored(self):
108        for f in (TESTFN2, TemporaryFile(), StringIO()):
109            self.zip_test(f, zipfile.ZIP_STORED)
110
111    def zip_open_test(self, f, compression):
112        self.make_test_archive(f, compression)
113
114        # Read the ZIP archive
115        with zipfile.ZipFile(f, "r", compression) as zipfp:
116            zipdata1 = []
117            zipopen1 = zipfp.open(TESTFN)
118            while True:
119                read_data = zipopen1.read(256)
120                if not read_data:
121                    break
122                zipdata1.append(read_data)
123
124            zipdata2 = []
125            zipopen2 = zipfp.open("another.name")
126            while True:
127                read_data = zipopen2.read(256)
128                if not read_data:
129                    break
130                zipdata2.append(read_data)
131
132            self.assertEqual(''.join(zipdata1), self.data)
133            self.assertEqual(''.join(zipdata2), self.data)
134
135    def test_open_stored(self):
136        for f in (TESTFN2, TemporaryFile(), StringIO()):
137            self.zip_open_test(f, zipfile.ZIP_STORED)
138
139    def test_open_via_zip_info(self):
140        # Create the ZIP archive
141        with zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED) as zipfp:
142            zipfp.writestr("name", "foo")
143            zipfp.writestr("name", "bar")
144
145        with zipfile.ZipFile(TESTFN2, "r") as zipfp:
146            infos = zipfp.infolist()
147            data = ""
148            for info in infos:
149                data += zipfp.open(info).read()
150            self.assertTrue(data == "foobar" or data == "barfoo")
151            data = ""
152            for info in infos:
153                data += zipfp.read(info)
154            self.assertTrue(data == "foobar" or data == "barfoo")
155
156    def zip_random_open_test(self, f, compression):
157        self.make_test_archive(f, compression)
158
159        # Read the ZIP archive
160        with zipfile.ZipFile(f, "r", compression) as zipfp:
161            zipdata1 = []
162            zipopen1 = zipfp.open(TESTFN)
163            while True:
164                read_data = zipopen1.read(randint(1, 1024))
165                if not read_data:
166                    break
167                zipdata1.append(read_data)
168
169            self.assertEqual(''.join(zipdata1), self.data)
170
171    def test_random_open_stored(self):
172        for f in (TESTFN2, TemporaryFile(), StringIO()):
173            self.zip_random_open_test(f, zipfile.ZIP_STORED)
174
175    def zip_readline_test(self, f, compression):
176        self.make_test_archive(f, compression)
177
178        # Read the ZIP archive
179        with zipfile.ZipFile(f, "r") as zipfp:
180            zipopen = zipfp.open(TESTFN)
181            for line in self.line_gen:
182                linedata = zipopen.readline()
183                self.assertEqual(linedata, line + '\n')
184
185    def zip_readlines_test(self, f, compression):
186        self.make_test_archive(f, compression)
187
188        # Read the ZIP archive
189        with zipfile.ZipFile(f, "r") as zipfp:
190            ziplines = zipfp.open(TESTFN).readlines()
191            for line, zipline in zip(self.line_gen, ziplines):
192                self.assertEqual(zipline, line + '\n')
193
194    def zip_iterlines_test(self, f, compression):
195        self.make_test_archive(f, compression)
196
197        # Read the ZIP archive
198        with zipfile.ZipFile(f, "r") as zipfp:
199            for line, zipline in zip(self.line_gen, zipfp.open(TESTFN)):
200                self.assertEqual(zipline, line + '\n')
201
202    def test_readline_stored(self):
203        for f in (TESTFN2, TemporaryFile(), StringIO()):
204            self.zip_readline_test(f, zipfile.ZIP_STORED)
205
206    def test_readlines_stored(self):
207        for f in (TESTFN2, TemporaryFile(), StringIO()):
208            self.zip_readlines_test(f, zipfile.ZIP_STORED)
209
210    def test_iterlines_stored(self):
211        for f in (TESTFN2, TemporaryFile(), StringIO()):
212            self.zip_iterlines_test(f, zipfile.ZIP_STORED)
213
214    @skipUnless(zlib, "requires zlib")
215    def test_deflated(self):
216        for f in (TESTFN2, TemporaryFile(), StringIO()):
217            self.zip_test(f, zipfile.ZIP_DEFLATED)
218
219    @skipUnless(zlib, "requires zlib")
220    def test_open_deflated(self):
221        for f in (TESTFN2, TemporaryFile(), StringIO()):
222            self.zip_open_test(f, zipfile.ZIP_DEFLATED)
223
224    @skipUnless(zlib, "requires zlib")
225    def test_random_open_deflated(self):
226        for f in (TESTFN2, TemporaryFile(), StringIO()):
227            self.zip_random_open_test(f, zipfile.ZIP_DEFLATED)
228
229    @skipUnless(zlib, "requires zlib")
230    def test_readline_deflated(self):
231        for f in (TESTFN2, TemporaryFile(), StringIO()):
232            self.zip_readline_test(f, zipfile.ZIP_DEFLATED)
233
234    @skipUnless(zlib, "requires zlib")
235    def test_readlines_deflated(self):
236        for f in (TESTFN2, TemporaryFile(), StringIO()):
237            self.zip_readlines_test(f, zipfile.ZIP_DEFLATED)
238
239    @skipUnless(zlib, "requires zlib")
240    def test_iterlines_deflated(self):
241        for f in (TESTFN2, TemporaryFile(), StringIO()):
242            self.zip_iterlines_test(f, zipfile.ZIP_DEFLATED)
243
244    @skipUnless(zlib, "requires zlib")
245    def test_low_compression(self):
246        """Check for cases where compressed data is larger than original."""
247        # Create the ZIP archive
248        with zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_DEFLATED) as zipfp:
249            zipfp.writestr("strfile", '12')
250
251        # Get an open object for strfile
252        with zipfile.ZipFile(TESTFN2, "r", zipfile.ZIP_DEFLATED) as zipfp:
253            openobj = zipfp.open("strfile")
254            self.assertEqual(openobj.read(1), '1')
255            self.assertEqual(openobj.read(1), '2')
256
257    def test_absolute_arcnames(self):
258        with zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED) as zipfp:
259            zipfp.write(TESTFN, "/absolute")
260
261        with zipfile.ZipFile(TESTFN2, "r", zipfile.ZIP_STORED) as zipfp:
262            self.assertEqual(zipfp.namelist(), ["absolute"])
263
264    def test_append_to_zip_file(self):
265        """Test appending to an existing zipfile."""
266        with zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED) as zipfp:
267            zipfp.write(TESTFN, TESTFN)
268
269        with zipfile.ZipFile(TESTFN2, "a", zipfile.ZIP_STORED) as zipfp:
270            zipfp.writestr("strfile", self.data)
271            self.assertEqual(zipfp.namelist(), [TESTFN, "strfile"])
272
273    def test_append_to_non_zip_file(self):
274        """Test appending to an existing file that is not a zipfile."""
275        # NOTE: this test fails if len(d) < 22 because of the first
276        # line "fpin.seek(-22, 2)" in _EndRecData
277        data = 'I am not a ZipFile!'*10
278        with open(TESTFN2, 'wb') as f:
279            f.write(data)
280
281        with zipfile.ZipFile(TESTFN2, "a", zipfile.ZIP_STORED) as zipfp:
282            zipfp.write(TESTFN, TESTFN)
283
284        with open(TESTFN2, 'rb') as f:
285            f.seek(len(data))
286            with zipfile.ZipFile(f, "r") as zipfp:
287                self.assertEqual(zipfp.namelist(), [TESTFN])
288
289    def test_write_default_name(self):
290        """Check that calling ZipFile.write without arcname specified
291        produces the expected result."""
292        with zipfile.ZipFile(TESTFN2, "w") as zipfp:
293            zipfp.write(TESTFN)
294            self.assertEqual(zipfp.read(TESTFN), open(TESTFN).read())
295
296    @skipUnless(zlib, "requires zlib")
297    def test_per_file_compression(self):
298        """Check that files within a Zip archive can have different
299        compression options."""
300        with zipfile.ZipFile(TESTFN2, "w") as zipfp:
301            zipfp.write(TESTFN, 'storeme', zipfile.ZIP_STORED)
302            zipfp.write(TESTFN, 'deflateme', zipfile.ZIP_DEFLATED)
303            sinfo = zipfp.getinfo('storeme')
304            dinfo = zipfp.getinfo('deflateme')
305            self.assertEqual(sinfo.compress_type, zipfile.ZIP_STORED)
306            self.assertEqual(dinfo.compress_type, zipfile.ZIP_DEFLATED)
307
308    def test_write_to_readonly(self):
309        """Check that trying to call write() on a readonly ZipFile object
310        raises a RuntimeError."""
311        with zipfile.ZipFile(TESTFN2, mode="w") as zipfp:
312            zipfp.writestr("somefile.txt", "bogus")
313
314        with zipfile.ZipFile(TESTFN2, mode="r") as zipfp:
315            self.assertRaises(RuntimeError, zipfp.write, TESTFN)
316
317    def test_extract(self):
318        with zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED) as zipfp:
319            for fpath, fdata in SMALL_TEST_DATA:
320                zipfp.writestr(fpath, fdata)
321
322        with zipfile.ZipFile(TESTFN2, "r") as zipfp:
323            for fpath, fdata in SMALL_TEST_DATA:
324                writtenfile = zipfp.extract(fpath)
325
326                # make sure it was written to the right place
327                if os.path.isabs(fpath):
328                    correctfile = os.path.join(os.getcwd(), fpath[1:])
329                else:
330                    correctfile = os.path.join(os.getcwd(), fpath)
331                correctfile = os.path.normpath(correctfile)
332
333                self.assertEqual(writtenfile, correctfile)
334
335                # make sure correct data is in correct file
336                self.assertEqual(fdata, open(writtenfile, "rb").read())
337                os.remove(writtenfile)
338
339        # remove the test file subdirectories
340        shutil.rmtree(os.path.join(os.getcwd(), 'ziptest2dir'))
341
342    def test_extract_all(self):
343        with zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED) as zipfp:
344            for fpath, fdata in SMALL_TEST_DATA:
345                zipfp.writestr(fpath, fdata)
346
347        with zipfile.ZipFile(TESTFN2, "r") as zipfp:
348            zipfp.extractall()
349            for fpath, fdata in SMALL_TEST_DATA:
350                if os.path.isabs(fpath):
351                    outfile = os.path.join(os.getcwd(), fpath[1:])
352                else:
353                    outfile = os.path.join(os.getcwd(), fpath)
354
355                self.assertEqual(fdata, open(outfile, "rb").read())
356                os.remove(outfile)
357
358        # remove the test file subdirectories
359        shutil.rmtree(os.path.join(os.getcwd(), 'ziptest2dir'))
360
361    def zip_test_writestr_permissions(self, f, compression):
362        # Make sure that writestr creates files with mode 0600,
363        # when it is passed a name rather than a ZipInfo instance.
364
365        self.make_test_archive(f, compression)
366        with zipfile.ZipFile(f, "r") as zipfp:
367            zinfo = zipfp.getinfo('strfile')
368            self.assertEqual(zinfo.external_attr, 0600 << 16)
369
370    def test_writestr_permissions(self):
371        for f in (TESTFN2, TemporaryFile(), StringIO()):
372            self.zip_test_writestr_permissions(f, zipfile.ZIP_STORED)
373
374    def test_close(self):
375        """Check that the zipfile is closed after the 'with' block."""
376        with zipfile.ZipFile(TESTFN2, "w") as zipfp:
377            for fpath, fdata in SMALL_TEST_DATA:
378                zipfp.writestr(fpath, fdata)
379                self.assertTrue(zipfp.fp is not None, 'zipfp is not open')
380        self.assertTrue(zipfp.fp is None, 'zipfp is not closed')
381
382        with zipfile.ZipFile(TESTFN2, "r") as zipfp:
383            self.assertTrue(zipfp.fp is not None, 'zipfp is not open')
384        self.assertTrue(zipfp.fp is None, 'zipfp is not closed')
385
386    def test_close_on_exception(self):
387        """Check that the zipfile is closed if an exception is raised in the
388        'with' block."""
389        with zipfile.ZipFile(TESTFN2, "w") as zipfp:
390            for fpath, fdata in SMALL_TEST_DATA:
391                zipfp.writestr(fpath, fdata)
392
393        try:
394            with zipfile.ZipFile(TESTFN2, "r") as zipfp2:
395                raise zipfile.BadZipfile()
396        except zipfile.BadZipfile:
397            self.assertTrue(zipfp2.fp is None, 'zipfp is not closed')
398
399    def tearDown(self):
400        unlink(TESTFN)
401        unlink(TESTFN2)
402
403
404class TestZip64InSmallFiles(unittest.TestCase):
405    # These tests test the ZIP64 functionality without using large files,
406    # see test_zipfile64 for proper tests.
407
408    def setUp(self):
409        self._limit = zipfile.ZIP64_LIMIT
410        zipfile.ZIP64_LIMIT = 5
411
412        line_gen = ("Test of zipfile line %d." % i
413                    for i in range(0, FIXEDTEST_SIZE))
414        self.data = '\n'.join(line_gen)
415
416        # Make a source file with some lines
417        with open(TESTFN, "wb") as fp:
418            fp.write(self.data)
419
420    def large_file_exception_test(self, f, compression):
421        with zipfile.ZipFile(f, "w", compression) as zipfp:
422            self.assertRaises(zipfile.LargeZipFile,
423                              zipfp.write, TESTFN, "another.name")
424
425    def large_file_exception_test2(self, f, compression):
426        with zipfile.ZipFile(f, "w", compression) as zipfp:
427            self.assertRaises(zipfile.LargeZipFile,
428                              zipfp.writestr, "another.name", self.data)
429
430    def test_large_file_exception(self):
431        for f in (TESTFN2, TemporaryFile(), StringIO()):
432            self.large_file_exception_test(f, zipfile.ZIP_STORED)
433            self.large_file_exception_test2(f, zipfile.ZIP_STORED)
434
435    def zip_test(self, f, compression):
436        # Create the ZIP archive
437        with zipfile.ZipFile(f, "w", compression, allowZip64=True) as zipfp:
438            zipfp.write(TESTFN, "another.name")
439            zipfp.write(TESTFN, TESTFN)
440            zipfp.writestr("strfile", self.data)
441
442        # Read the ZIP archive
443        with zipfile.ZipFile(f, "r", compression) as zipfp:
444            self.assertEqual(zipfp.read(TESTFN), self.data)
445            self.assertEqual(zipfp.read("another.name"), self.data)
446            self.assertEqual(zipfp.read("strfile"), self.data)
447
448            # Print the ZIP directory
449            fp = StringIO()
450            stdout = sys.stdout
451            try:
452                sys.stdout = fp
453                zipfp.printdir()
454            finally:
455                sys.stdout = stdout
456
457            directory = fp.getvalue()
458            lines = directory.splitlines()
459            self.assertEqual(len(lines), 4) # Number of files + header
460
461            self.assertTrue('File Name' in lines[0])
462            self.assertTrue('Modified' in lines[0])
463            self.assertTrue('Size' in lines[0])
464
465            fn, date, time_, size = lines[1].split()
466            self.assertEqual(fn, 'another.name')
467            self.assertTrue(time.strptime(date, '%Y-%m-%d'))
468            self.assertTrue(time.strptime(time_, '%H:%M:%S'))
469            self.assertEqual(size, str(len(self.data)))
470
471            # Check the namelist
472            names = zipfp.namelist()
473            self.assertEqual(len(names), 3)
474            self.assertTrue(TESTFN in names)
475            self.assertTrue("another.name" in names)
476            self.assertTrue("strfile" in names)
477
478            # Check infolist
479            infos = zipfp.infolist()
480            names = [i.filename for i in infos]
481            self.assertEqual(len(names), 3)
482            self.assertTrue(TESTFN in names)
483            self.assertTrue("another.name" in names)
484            self.assertTrue("strfile" in names)
485            for i in infos:
486                self.assertEqual(i.file_size, len(self.data))
487
488            # check getinfo
489            for nm in (TESTFN, "another.name", "strfile"):
490                info = zipfp.getinfo(nm)
491                self.assertEqual(info.filename, nm)
492                self.assertEqual(info.file_size, len(self.data))
493
494            # Check that testzip doesn't raise an exception
495            zipfp.testzip()
496
497    def test_stored(self):
498        for f in (TESTFN2, TemporaryFile(), StringIO()):
499            self.zip_test(f, zipfile.ZIP_STORED)
500
501    @skipUnless(zlib, "requires zlib")
502    def test_deflated(self):
503        for f in (TESTFN2, TemporaryFile(), StringIO()):
504            self.zip_test(f, zipfile.ZIP_DEFLATED)
505
506    def test_absolute_arcnames(self):
507        with zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED,
508                             allowZip64=True) as zipfp:
509            zipfp.write(TESTFN, "/absolute")
510
511        with zipfile.ZipFile(TESTFN2, "r", zipfile.ZIP_STORED) as zipfp:
512            self.assertEqual(zipfp.namelist(), ["absolute"])
513
514    def tearDown(self):
515        zipfile.ZIP64_LIMIT = self._limit
516        unlink(TESTFN)
517        unlink(TESTFN2)
518
519
520class PyZipFileTests(unittest.TestCase):
521    def test_write_pyfile(self):
522        with zipfile.PyZipFile(TemporaryFile(), "w") as zipfp:
523            fn = __file__
524            if fn.endswith('.pyc') or fn.endswith('.pyo'):
525                fn = fn[:-1]
526
527            zipfp.writepy(fn)
528
529            bn = os.path.basename(fn)
530            self.assertTrue(bn not in zipfp.namelist())
531            self.assertTrue(bn + 'o' in zipfp.namelist() or
532                            bn + 'c' in zipfp.namelist())
533
534        with zipfile.PyZipFile(TemporaryFile(), "w") as zipfp:
535            fn = __file__
536            if fn.endswith(('.pyc', '.pyo')):
537                fn = fn[:-1]
538
539            zipfp.writepy(fn, "testpackage")
540
541            bn = "%s/%s" % ("testpackage", os.path.basename(fn))
542            self.assertTrue(bn not in zipfp.namelist())
543            self.assertTrue(bn + 'o' in zipfp.namelist() or
544                            bn + 'c' in zipfp.namelist())
545
546    def test_write_python_package(self):
547        import email
548        packagedir = os.path.dirname(email.__file__)
549
550        with zipfile.PyZipFile(TemporaryFile(), "w") as zipfp:
551            zipfp.writepy(packagedir)
552
553            # Check for a couple of modules at different levels of the
554            # hierarchy
555            names = zipfp.namelist()
556            self.assertTrue('email/__init__.pyo' in names or
557                            'email/__init__.pyc' in names)
558            self.assertTrue('email/mime/text.pyo' in names or
559                            'email/mime/text.pyc' in names)
560
561    def test_write_python_directory(self):
562        os.mkdir(TESTFN2)
563        try:
564            with open(os.path.join(TESTFN2, "mod1.py"), "w") as fp:
565                fp.write("print 42\n")
566
567            with open(os.path.join(TESTFN2, "mod2.py"), "w") as fp:
568                fp.write("print 42 * 42\n")
569
570            with open(os.path.join(TESTFN2, "mod2.txt"), "w") as fp:
571                fp.write("bla bla bla\n")
572
573            zipfp  = zipfile.PyZipFile(TemporaryFile(), "w")
574            zipfp.writepy(TESTFN2)
575
576            names = zipfp.namelist()
577            self.assertTrue('mod1.pyc' in names or 'mod1.pyo' in names)
578            self.assertTrue('mod2.pyc' in names or 'mod2.pyo' in names)
579            self.assertTrue('mod2.txt' not in names)
580
581        finally:
582            shutil.rmtree(TESTFN2)
583
584    def test_write_non_pyfile(self):
585        with zipfile.PyZipFile(TemporaryFile(), "w") as zipfp:
586            open(TESTFN, 'w').write('most definitely not a python file')
587            self.assertRaises(RuntimeError, zipfp.writepy, TESTFN)
588            os.remove(TESTFN)
589
590
591class OtherTests(unittest.TestCase):
592    def test_unicode_filenames(self):
593        with zipfile.ZipFile(TESTFN, "w") as zf:
594            zf.writestr(u"foo.txt", "Test for unicode filename")
595            zf.writestr(u"\xf6.txt", "Test for unicode filename")
596            self.assertTrue(isinstance(zf.infolist()[0].filename, unicode))
597
598        with zipfile.ZipFile(TESTFN, "r") as zf:
599            self.assertEqual(zf.filelist[0].filename, "foo.txt")
600            self.assertEqual(zf.filelist[1].filename, u"\xf6.txt")
601
602    def test_create_non_existent_file_for_append(self):
603        if os.path.exists(TESTFN):
604            os.unlink(TESTFN)
605
606        filename = 'testfile.txt'
607        content = 'hello, world. this is some content.'
608
609        try:
610            with zipfile.ZipFile(TESTFN, 'a') as zf:
611                zf.writestr(filename, content)
612        except IOError:
613            self.fail('Could not append data to a non-existent zip file.')
614
615        self.assertTrue(os.path.exists(TESTFN))
616
617        with zipfile.ZipFile(TESTFN, 'r') as zf:
618            self.assertEqual(zf.read(filename), content)
619
620    def test_close_erroneous_file(self):
621        # This test checks that the ZipFile constructor closes the file object
622        # it opens if there's an error in the file.  If it doesn't, the
623        # traceback holds a reference to the ZipFile object and, indirectly,
624        # the file object.
625        # On Windows, this causes the os.unlink() call to fail because the
626        # underlying file is still open.  This is SF bug #412214.
627        #
628        with open(TESTFN, "w") as fp:
629            fp.write("this is not a legal zip file\n")
630        try:
631            zf = zipfile.ZipFile(TESTFN)
632        except zipfile.BadZipfile:
633            pass
634
635    def test_is_zip_erroneous_file(self):
636        """Check that is_zipfile() correctly identifies non-zip files."""
637        # - passing a filename
638        with open(TESTFN, "w") as fp:
639            fp.write("this is not a legal zip file\n")
640        chk = zipfile.is_zipfile(TESTFN)
641        self.assertFalse(chk)
642        # - passing a file object
643        with open(TESTFN, "rb") as fp:
644            chk = zipfile.is_zipfile(fp)
645            self.assertTrue(not chk)
646        # - passing a file-like object
647        fp = StringIO()
648        fp.write("this is not a legal zip file\n")
649        chk = zipfile.is_zipfile(fp)
650        self.assertTrue(not chk)
651        fp.seek(0, 0)
652        chk = zipfile.is_zipfile(fp)
653        self.assertTrue(not chk)
654
655    def test_is_zip_valid_file(self):
656        """Check that is_zipfile() correctly identifies zip files."""
657        # - passing a filename
658        with zipfile.ZipFile(TESTFN, mode="w") as zipf:
659            zipf.writestr("foo.txt", "O, for a Muse of Fire!")
660        chk = zipfile.is_zipfile(TESTFN)
661        self.assertTrue(chk)
662        # - passing a file object
663        with open(TESTFN, "rb") as fp:
664            chk = zipfile.is_zipfile(fp)
665            self.assertTrue(chk)
666            fp.seek(0, 0)
667            zip_contents = fp.read()
668        # - passing a file-like object
669        fp = StringIO()
670        fp.write(zip_contents)
671        chk = zipfile.is_zipfile(fp)
672        self.assertTrue(chk)
673        fp.seek(0, 0)
674        chk = zipfile.is_zipfile(fp)
675        self.assertTrue(chk)
676
677    def test_non_existent_file_raises_IOError(self):
678        # make sure we don't raise an AttributeError when a partially-constructed
679        # ZipFile instance is finalized; this tests for regression on SF tracker
680        # bug #403871.
681
682        # The bug we're testing for caused an AttributeError to be raised
683        # when a ZipFile instance was created for a file that did not
684        # exist; the .fp member was not initialized but was needed by the
685        # __del__() method.  Since the AttributeError is in the __del__(),
686        # it is ignored, but the user should be sufficiently annoyed by
687        # the message on the output that regression will be noticed
688        # quickly.
689        self.assertRaises(IOError, zipfile.ZipFile, TESTFN)
690
691    def test_empty_file_raises_BadZipFile(self):
692        f = open(TESTFN, 'w')
693        f.close()
694        self.assertRaises(zipfile.BadZipfile, zipfile.ZipFile, TESTFN)
695
696        with open(TESTFN, 'w') as fp:
697            fp.write("short file")
698        self.assertRaises(zipfile.BadZipfile, zipfile.ZipFile, TESTFN)
699
700    def test_closed_zip_raises_RuntimeError(self):
701        """Verify that testzip() doesn't swallow inappropriate exceptions."""
702        data = StringIO()
703        with zipfile.ZipFile(data, mode="w") as zipf:
704            zipf.writestr("foo.txt", "O, for a Muse of Fire!")
705
706        # This is correct; calling .read on a closed ZipFile should throw
707        # a RuntimeError, and so should calling .testzip.  An earlier
708        # version of .testzip would swallow this exception (and any other)
709        # and report that the first file in the archive was corrupt.
710        self.assertRaises(RuntimeError, zipf.read, "foo.txt")
711        self.assertRaises(RuntimeError, zipf.open, "foo.txt")
712        self.assertRaises(RuntimeError, zipf.testzip)
713        self.assertRaises(RuntimeError, zipf.writestr, "bogus.txt", "bogus")
714        open(TESTFN, 'w').write('zipfile test data')
715        self.assertRaises(RuntimeError, zipf.write, TESTFN)
716
717    def test_bad_constructor_mode(self):
718        """Check that bad modes passed to ZipFile constructor are caught."""
719        self.assertRaises(RuntimeError, zipfile.ZipFile, TESTFN, "q")
720
721    def test_bad_open_mode(self):
722        """Check that bad modes passed to ZipFile.open are caught."""
723        with zipfile.ZipFile(TESTFN, mode="w") as zipf:
724            zipf.writestr("foo.txt", "O, for a Muse of Fire!")
725
726        with zipfile.ZipFile(TESTFN, mode="r") as zipf:
727        # read the data to make sure the file is there
728            zipf.read("foo.txt")
729            self.assertRaises(RuntimeError, zipf.open, "foo.txt", "q")
730
731    def test_read0(self):
732        """Check that calling read(0) on a ZipExtFile object returns an empty
733        string and doesn't advance file pointer."""
734        with zipfile.ZipFile(TESTFN, mode="w") as zipf:
735            zipf.writestr("foo.txt", "O, for a Muse of Fire!")
736            # read the data to make sure the file is there
737            f = zipf.open("foo.txt")
738            for i in xrange(FIXEDTEST_SIZE):
739                self.assertEqual(f.read(0), '')
740
741            self.assertEqual(f.read(), "O, for a Muse of Fire!")
742
743    def test_open_non_existent_item(self):
744        """Check that attempting to call open() for an item that doesn't
745        exist in the archive raises a RuntimeError."""
746        with zipfile.ZipFile(TESTFN, mode="w") as zipf:
747            self.assertRaises(KeyError, zipf.open, "foo.txt", "r")
748
749    def test_bad_compression_mode(self):
750        """Check that bad compression methods passed to ZipFile.open are
751        caught."""
752        self.assertRaises(RuntimeError, zipfile.ZipFile, TESTFN, "w", -1)
753
754    def test_null_byte_in_filename(self):
755        """Check that a filename containing a null byte is properly
756        terminated."""
757        with zipfile.ZipFile(TESTFN, mode="w") as zipf:
758            zipf.writestr("foo.txt\x00qqq", "O, for a Muse of Fire!")
759            self.assertEqual(zipf.namelist(), ['foo.txt'])
760
761    def test_struct_sizes(self):
762        """Check that ZIP internal structure sizes are calculated correctly."""
763        self.assertEqual(zipfile.sizeEndCentDir, 22)
764        self.assertEqual(zipfile.sizeCentralDir, 46)
765        self.assertEqual(zipfile.sizeEndCentDir64, 56)
766        self.assertEqual(zipfile.sizeEndCentDir64Locator, 20)
767
768    def test_comments(self):
769        """Check that comments on the archive are handled properly."""
770
771        # check default comment is empty
772        with zipfile.ZipFile(TESTFN, mode="w") as zipf:
773            self.assertEqual(zipf.comment, '')
774            zipf.writestr("foo.txt", "O, for a Muse of Fire!")
775
776        with zipfile.ZipFile(TESTFN, mode="r") as zipf:
777            self.assertEqual(zipf.comment, '')
778
779        # check a simple short comment
780        comment = 'Bravely taking to his feet, he beat a very brave retreat.'
781        with zipfile.ZipFile(TESTFN, mode="w") as zipf:
782            zipf.comment = comment
783            zipf.writestr("foo.txt", "O, for a Muse of Fire!")
784        with zipfile.ZipFile(TESTFN, mode="r") as zipf:
785            self.assertEqual(zipf.comment, comment)
786
787        # check a comment of max length
788        comment2 = ''.join(['%d' % (i**3 % 10) for i in xrange((1 << 16)-1)])
789        with zipfile.ZipFile(TESTFN, mode="w") as zipf:
790            zipf.comment = comment2
791            zipf.writestr("foo.txt", "O, for a Muse of Fire!")
792
793        with zipfile.ZipFile(TESTFN, mode="r") as zipf:
794            self.assertEqual(zipf.comment, comment2)
795
796        # check a comment that is too long is truncated
797        with zipfile.ZipFile(TESTFN, mode="w") as zipf:
798            zipf.comment = comment2 + 'oops'
799            zipf.writestr("foo.txt", "O, for a Muse of Fire!")
800        with zipfile.ZipFile(TESTFN, mode="r") as zipf:
801            self.assertEqual(zipf.comment, comment2)
802
803    def tearDown(self):
804        unlink(TESTFN)
805        unlink(TESTFN2)
806
807
808class DecryptionTests(unittest.TestCase):
809    """Check that ZIP decryption works. Since the library does not
810    support encryption at the moment, we use a pre-generated encrypted
811    ZIP file."""
812
813    data = (
814    'PK\x03\x04\x14\x00\x01\x00\x00\x00n\x92i.#y\xef?&\x00\x00\x00\x1a\x00'
815    '\x00\x00\x08\x00\x00\x00test.txt\xfa\x10\xa0gly|\xfa-\xc5\xc0=\xf9y'
816    '\x18\xe0\xa8r\xb3Z}Lg\xbc\xae\xf9|\x9b\x19\xe4\x8b\xba\xbb)\x8c\xb0\xdbl'
817    'PK\x01\x02\x14\x00\x14\x00\x01\x00\x00\x00n\x92i.#y\xef?&\x00\x00\x00'
818    '\x1a\x00\x00\x00\x08\x00\x00\x00\x00\x00\x00\x00\x01\x00 \x00\xb6\x81'
819    '\x00\x00\x00\x00test.txtPK\x05\x06\x00\x00\x00\x00\x01\x00\x01\x006\x00'
820    '\x00\x00L\x00\x00\x00\x00\x00' )
821    data2 = (
822    'PK\x03\x04\x14\x00\t\x00\x08\x00\xcf}38xu\xaa\xb2\x14\x00\x00\x00\x00\x02'
823    '\x00\x00\x04\x00\x15\x00zeroUT\t\x00\x03\xd6\x8b\x92G\xda\x8b\x92GUx\x04'
824    '\x00\xe8\x03\xe8\x03\xc7<M\xb5a\xceX\xa3Y&\x8b{oE\xd7\x9d\x8c\x98\x02\xc0'
825    'PK\x07\x08xu\xaa\xb2\x14\x00\x00\x00\x00\x02\x00\x00PK\x01\x02\x17\x03'
826    '\x14\x00\t\x00\x08\x00\xcf}38xu\xaa\xb2\x14\x00\x00\x00\x00\x02\x00\x00'
827    '\x04\x00\r\x00\x00\x00\x00\x00\x00\x00\x00\x00\xa4\x81\x00\x00\x00\x00ze'
828    'roUT\x05\x00\x03\xd6\x8b\x92GUx\x00\x00PK\x05\x06\x00\x00\x00\x00\x01'
829    '\x00\x01\x00?\x00\x00\x00[\x00\x00\x00\x00\x00' )
830
831    plain = 'zipfile.py encryption test'
832    plain2 = '\x00'*512
833
834    def setUp(self):
835        with open(TESTFN, "wb") as fp:
836            fp.write(self.data)
837        self.zip = zipfile.ZipFile(TESTFN, "r")
838        with open(TESTFN2, "wb") as fp:
839            fp.write(self.data2)
840        self.zip2 = zipfile.ZipFile(TESTFN2, "r")
841
842    def tearDown(self):
843        self.zip.close()
844        os.unlink(TESTFN)
845        self.zip2.close()
846        os.unlink(TESTFN2)
847
848    def test_no_password(self):
849        # Reading the encrypted file without password
850        # must generate a RunTime exception
851        self.assertRaises(RuntimeError, self.zip.read, "test.txt")
852        self.assertRaises(RuntimeError, self.zip2.read, "zero")
853
854    def test_bad_password(self):
855        self.zip.setpassword("perl")
856        self.assertRaises(RuntimeError, self.zip.read, "test.txt")
857        self.zip2.setpassword("perl")
858        self.assertRaises(RuntimeError, self.zip2.read, "zero")
859
860    @skipUnless(zlib, "requires zlib")
861    def test_good_password(self):
862        self.zip.setpassword("python")
863        self.assertEqual(self.zip.read("test.txt"), self.plain)
864        self.zip2.setpassword("12345")
865        self.assertEqual(self.zip2.read("zero"), self.plain2)
866
867
868class TestsWithRandomBinaryFiles(unittest.TestCase):
869    def setUp(self):
870        datacount = randint(16, 64)*1024 + randint(1, 1024)
871        self.data = ''.join(struct.pack('<f', random()*randint(-1000, 1000))
872                                        for i in xrange(datacount))
873
874        # Make a source file with some lines
875        with open(TESTFN, "wb") as fp:
876            fp.write(self.data)
877
878    def tearDown(self):
879        unlink(TESTFN)
880        unlink(TESTFN2)
881
882    def make_test_archive(self, f, compression):
883        # Create the ZIP archive
884        with zipfile.ZipFile(f, "w", compression) as zipfp:
885            zipfp.write(TESTFN, "another.name")
886            zipfp.write(TESTFN, TESTFN)
887
888    def zip_test(self, f, compression):
889        self.make_test_archive(f, compression)
890
891        # Read the ZIP archive
892        with zipfile.ZipFile(f, "r", compression) as zipfp:
893            testdata = zipfp.read(TESTFN)
894            self.assertEqual(len(testdata), len(self.data))
895            self.assertEqual(testdata, self.data)
896            self.assertEqual(zipfp.read("another.name"), self.data)
897
898    def test_stored(self):
899        for f in (TESTFN2, TemporaryFile(), StringIO()):
900            self.zip_test(f, zipfile.ZIP_STORED)
901
902    def zip_open_test(self, f, compression):
903        self.make_test_archive(f, compression)
904
905        # Read the ZIP archive
906        with zipfile.ZipFile(f, "r", compression) as zipfp:
907            zipdata1 = []
908            zipopen1 = zipfp.open(TESTFN)
909            while True:
910                read_data = zipopen1.read(256)
911                if not read_data:
912                    break
913                zipdata1.append(read_data)
914
915            zipdata2 = []
916            zipopen2 = zipfp.open("another.name")
917            while True:
918                read_data = zipopen2.read(256)
919                if not read_data:
920                    break
921                zipdata2.append(read_data)
922
923            testdata1 = ''.join(zipdata1)
924            self.assertEqual(len(testdata1), len(self.data))
925            self.assertEqual(testdata1, self.data)
926
927            testdata2 = ''.join(zipdata2)
928            self.assertEqual(len(testdata2), len(self.data))
929            self.assertEqual(testdata2, self.data)
930
931    def test_open_stored(self):
932        for f in (TESTFN2, TemporaryFile(), StringIO()):
933            self.zip_open_test(f, zipfile.ZIP_STORED)
934
935    def zip_random_open_test(self, f, compression):
936        self.make_test_archive(f, compression)
937
938        # Read the ZIP archive
939        with zipfile.ZipFile(f, "r", compression) as zipfp:
940            zipdata1 = []
941            zipopen1 = zipfp.open(TESTFN)
942            while True:
943                read_data = zipopen1.read(randint(1, 1024))
944                if not read_data:
945                    break
946                zipdata1.append(read_data)
947
948            testdata = ''.join(zipdata1)
949            self.assertEqual(len(testdata), len(self.data))
950            self.assertEqual(testdata, self.data)
951
952    def test_random_open_stored(self):
953        for f in (TESTFN2, TemporaryFile(), StringIO()):
954            self.zip_random_open_test(f, zipfile.ZIP_STORED)
955
956
957@skipUnless(zlib, "requires zlib")
958class TestsWithMultipleOpens(unittest.TestCase):
959    def setUp(self):
960        # Create the ZIP archive
961        with zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_DEFLATED) as zipfp:
962            zipfp.writestr('ones', '1'*FIXEDTEST_SIZE)
963            zipfp.writestr('twos', '2'*FIXEDTEST_SIZE)
964
965    def test_same_file(self):
966        # Verify that (when the ZipFile is in control of creating file objects)
967        # multiple open() calls can be made without interfering with each other.
968        with zipfile.ZipFile(TESTFN2, mode="r") as zipf:
969            zopen1 = zipf.open('ones')
970            zopen2 = zipf.open('ones')
971            data1 = zopen1.read(500)
972            data2 = zopen2.read(500)
973            data1 += zopen1.read(500)
974            data2 += zopen2.read(500)
975            self.assertEqual(data1, data2)
976
977    def test_different_file(self):
978        # Verify that (when the ZipFile is in control of creating file objects)
979        # multiple open() calls can be made without interfering with each other.
980        with zipfile.ZipFile(TESTFN2, mode="r") as zipf:
981            zopen1 = zipf.open('ones')
982            zopen2 = zipf.open('twos')
983            data1 = zopen1.read(500)
984            data2 = zopen2.read(500)
985            data1 += zopen1.read(500)
986            data2 += zopen2.read(500)
987            self.assertEqual(data1, '1'*FIXEDTEST_SIZE)
988            self.assertEqual(data2, '2'*FIXEDTEST_SIZE)
989
990    def test_interleaved(self):
991        # Verify that (when the ZipFile is in control of creating file objects)
992        # multiple open() calls can be made without interfering with each other.
993        with zipfile.ZipFile(TESTFN2, mode="r") as zipf:
994            zopen1 = zipf.open('ones')
995            data1 = zopen1.read(500)
996            zopen2 = zipf.open('twos')
997            data2 = zopen2.read(500)
998            data1 += zopen1.read(500)
999            data2 += zopen2.read(500)
1000            self.assertEqual(data1, '1'*FIXEDTEST_SIZE)
1001            self.assertEqual(data2, '2'*FIXEDTEST_SIZE)
1002
1003    def tearDown(self):
1004        unlink(TESTFN2)
1005
1006
1007class TestWithDirectory(unittest.TestCase):
1008    def setUp(self):
1009        os.mkdir(TESTFN2)
1010
1011    def test_extract_dir(self):
1012        with zipfile.ZipFile(findfile("zipdir.zip")) as zipf:
1013            zipf.extractall(TESTFN2)
1014        self.assertTrue(os.path.isdir(os.path.join(TESTFN2, "a")))
1015        self.assertTrue(os.path.isdir(os.path.join(TESTFN2, "a", "b")))
1016        self.assertTrue(os.path.exists(os.path.join(TESTFN2, "a", "b", "c")))
1017
1018    def test_bug_6050(self):
1019        # Extraction should succeed if directories already exist
1020        os.mkdir(os.path.join(TESTFN2, "a"))
1021        self.test_extract_dir()
1022
1023    def test_store_dir(self):
1024        os.mkdir(os.path.join(TESTFN2, "x"))
1025        zipf = zipfile.ZipFile(TESTFN, "w")
1026        zipf.write(os.path.join(TESTFN2, "x"), "x")
1027        self.assertTrue(zipf.filelist[0].filename.endswith("x/"))
1028
1029    def tearDown(self):
1030        shutil.rmtree(TESTFN2)
1031        if os.path.exists(TESTFN):
1032            unlink(TESTFN)
1033
1034
1035class UniversalNewlineTests(unittest.TestCase):
1036    def setUp(self):
1037        self.line_gen = ["Test of zipfile line %d." % i
1038                         for i in xrange(FIXEDTEST_SIZE)]
1039        self.seps = ('\r', '\r\n', '\n')
1040        self.arcdata, self.arcfiles = {}, {}
1041        for n, s in enumerate(self.seps):
1042            self.arcdata[s] = s.join(self.line_gen) + s
1043            self.arcfiles[s] = '%s-%d' % (TESTFN, n)
1044            open(self.arcfiles[s], "wb").write(self.arcdata[s])
1045
1046    def make_test_archive(self, f, compression):
1047        # Create the ZIP archive
1048        with zipfile.ZipFile(f, "w", compression) as zipfp:
1049            for fn in self.arcfiles.values():
1050                zipfp.write(fn, fn)
1051
1052    def read_test(self, f, compression):
1053        self.make_test_archive(f, compression)
1054
1055        # Read the ZIP archive
1056        with zipfile.ZipFile(f, "r") as zipfp:
1057            for sep, fn in self.arcfiles.items():
1058                zipdata = zipfp.open(fn, "rU").read()
1059                self.assertEqual(self.arcdata[sep], zipdata)
1060
1061    def readline_test(self, f, compression):
1062        self.make_test_archive(f, compression)
1063
1064        # Read the ZIP archive
1065        with zipfile.ZipFile(f, "r") as zipfp:
1066            for sep, fn in self.arcfiles.items():
1067                zipopen = zipfp.open(fn, "rU")
1068                for line in self.line_gen:
1069                    linedata = zipopen.readline()
1070                    self.assertEqual(linedata, line + '\n')
1071
1072    def readlines_test(self, f, compression):
1073        self.make_test_archive(f, compression)
1074
1075        # Read the ZIP archive
1076        with zipfile.ZipFile(f, "r") as zipfp:
1077            for sep, fn in self.arcfiles.items():
1078                ziplines = zipfp.open(fn, "rU").readlines()
1079                for line, zipline in zip(self.line_gen, ziplines):
1080                    self.assertEqual(zipline, line + '\n')
1081
1082    def iterlines_test(self, f, compression):
1083        self.make_test_archive(f, compression)
1084
1085        # Read the ZIP archive
1086        with zipfile.ZipFile(f, "r") as zipfp:
1087            for sep, fn in self.arcfiles.items():
1088                for line, zipline in zip(self.line_gen, zipfp.open(fn, "rU")):
1089                    self.assertEqual(zipline, line + '\n')
1090
1091    def test_read_stored(self):
1092        for f in (TESTFN2, TemporaryFile(), StringIO()):
1093            self.read_test(f, zipfile.ZIP_STORED)
1094
1095    def test_readline_stored(self):
1096        for f in (TESTFN2, TemporaryFile(), StringIO()):
1097            self.readline_test(f, zipfile.ZIP_STORED)
1098
1099    def test_readlines_stored(self):
1100        for f in (TESTFN2, TemporaryFile(), StringIO()):
1101            self.readlines_test(f, zipfile.ZIP_STORED)
1102
1103    def test_iterlines_stored(self):
1104        for f in (TESTFN2, TemporaryFile(), StringIO()):
1105            self.iterlines_test(f, zipfile.ZIP_STORED)
1106
1107    @skipUnless(zlib, "requires zlib")
1108    def test_read_deflated(self):
1109        for f in (TESTFN2, TemporaryFile(), StringIO()):
1110            self.read_test(f, zipfile.ZIP_DEFLATED)
1111
1112    @skipUnless(zlib, "requires zlib")
1113    def test_readline_deflated(self):
1114        for f in (TESTFN2, TemporaryFile(), StringIO()):
1115            self.readline_test(f, zipfile.ZIP_DEFLATED)
1116
1117    @skipUnless(zlib, "requires zlib")
1118    def test_readlines_deflated(self):
1119        for f in (TESTFN2, TemporaryFile(), StringIO()):
1120            self.readlines_test(f, zipfile.ZIP_DEFLATED)
1121
1122    @skipUnless(zlib, "requires zlib")
1123    def test_iterlines_deflated(self):
1124        for f in (TESTFN2, TemporaryFile(), StringIO()):
1125            self.iterlines_test(f, zipfile.ZIP_DEFLATED)
1126
1127    def tearDown(self):
1128        for sep, fn in self.arcfiles.items():
1129            os.remove(fn)
1130        unlink(TESTFN)
1131        unlink(TESTFN2)
1132
1133
1134def test_main():
1135    run_unittest(TestsWithSourceFile, TestZip64InSmallFiles, OtherTests,
1136                 PyZipFileTests, DecryptionTests, TestsWithMultipleOpens,
1137                 TestWithDirectory, UniversalNewlineTests,
1138                 TestsWithRandomBinaryFiles)
1139
1140if __name__ == "__main__":
1141    test_main()
1142