1 # We can test part of the module without zlib.

     2 try:

     3     import zlib

     4 except ImportError:

     5     zlib = None

     6 

     7 import zipfile, os, unittest, sys, shutil

     8 

     9 from StringIO import StringIO

    10 from tempfile import TemporaryFile

    11 

    12 from test.test_support import TESTFN, run_unittest

    13 

    14 TESTFN2 = TESTFN + "2"

    15 

    16 class TestsWithSourceFile(unittest.TestCase):

    17     def setUp(self):

    18         line_gen = ("Test of zipfile line %d." % i for i in range(0, 1000))

    19         self.data = '\n'.join(line_gen)

    20 

    21         # Make a source file with some lines

    22         fp = open(TESTFN, "wb")

    23         fp.write(self.data)

    24         fp.close()

    25 

    26     def zipTest(self, f, compression):

    27         # Create the ZIP archive

    28         zipfp = zipfile.ZipFile(f, "w", compression)

    29         zipfp.write(TESTFN, "another"+os.extsep+"name")

    30         zipfp.write(TESTFN, TESTFN)

    31         zipfp.writestr("strfile", self.data)

    32         zipfp.close()

    33 

    34         # Read the ZIP archive

    35         zipfp = zipfile.ZipFile(f, "r", compression)

    36         self.assertEqual(zipfp.read(TESTFN), self.data)

    37         self.assertEqual(zipfp.read("another"+os.extsep+"name"), self.data)

    38         self.assertEqual(zipfp.read("strfile"), self.data)

    39 

    40         # Print the ZIP directory

    41         fp = StringIO()

    42         stdout = sys.stdout

    43         try:

    44             sys.stdout = fp

    45 

    46             zipfp.printdir()

    47         finally:

    48             sys.stdout = stdout

    49 

    50         directory = fp.getvalue()

    51         lines = directory.splitlines()

    52         self.assertEquals(len(lines), 4) # Number of files + header

    53 

    54         self.assert_('File Name' in lines[0])

    55         self.assert_('Modified' in lines[0])

    56         self.assert_('Size' in lines[0])

    57 

    58         fn, date, time, size = lines[1].split()

    59         self.assertEquals(fn, 'another.name')

    60         # XXX: timestamp is not tested

    61         self.assertEquals(size, str(len(self.data)))

    62 

    63         # Check the namelist

    64         names = zipfp.namelist()

    65         self.assertEquals(len(names), 3)

    66         self.assert_(TESTFN in names)

    67         self.assert_("another"+os.extsep+"name" in names)

    68         self.assert_("strfile" in names)

    69 

    70         # Check infolist

    71         infos = zipfp.infolist()

    72         names = [ i.filename for i in infos ]

    73         self.assertEquals(len(names), 3)

    74         self.assert_(TESTFN in names)

    75         self.assert_("another"+os.extsep+"name" in names)

    76         self.assert_("strfile" in names)

    77         for i in infos:

    78             self.assertEquals(i.file_size, len(self.data))

    79 

    80         # check getinfo

    81         for nm in (TESTFN, "another"+os.extsep+"name", "strfile"):

    82             info = zipfp.getinfo(nm)

    83             self.assertEquals(info.filename, nm)

    84             self.assertEquals(info.file_size, len(self.data))

    85 

    86         # Check that testzip doesn't raise an exception

    87         zipfp.testzip()

    88 

    89 

    90         zipfp.close()

    91 

    92 

    93 

    94 

    95     def testStored(self):

    96         for f in (TESTFN2, TemporaryFile(), StringIO()):

    97             self.zipTest(f, zipfile.ZIP_STORED)

    98 

    99     if zlib:

   100         def testDeflated(self):

   101             for f in (TESTFN2, TemporaryFile(), StringIO()):

   102                 self.zipTest(f, zipfile.ZIP_DEFLATED)

   103 

   104     def testAbsoluteArcnames(self):

   105         zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED)

   106         zipfp.write(TESTFN, "/absolute")

   107         zipfp.close()

   108 

   109         zipfp = zipfile.ZipFile(TESTFN2, "r", zipfile.ZIP_STORED)

   110         self.assertEqual(zipfp.namelist(), ["absolute"])

   111         zipfp.close()

   112 

   113 

   114     def tearDown(self):

   115         os.remove(TESTFN)

   116         os.remove(TESTFN2)

   117 

   118 class TestZip64InSmallFiles(unittest.TestCase):

   119     # These tests test the ZIP64 functionality without using large files,

   120     # see test_zipfile64 for proper tests.

   121 

   122     def setUp(self):

   123         self._limit = zipfile.ZIP64_LIMIT

   124         zipfile.ZIP64_LIMIT = 5

   125 

   126         line_gen = ("Test of zipfile line %d." % i for i in range(0, 1000))

   127         self.data = '\n'.join(line_gen)

   128 

   129         # Make a source file with some lines

   130         fp = open(TESTFN, "wb")

   131         fp.write(self.data)

   132         fp.close()

   133 

   134     def largeFileExceptionTest(self, f, compression):

   135         zipfp = zipfile.ZipFile(f, "w", compression)

   136         self.assertRaises(zipfile.LargeZipFile,

   137                 zipfp.write, TESTFN, "another"+os.extsep+"name")

   138         zipfp.close()

   139 

   140     def largeFileExceptionTest2(self, f, compression):

   141         zipfp = zipfile.ZipFile(f, "w", compression)

   142         self.assertRaises(zipfile.LargeZipFile,

   143                 zipfp.writestr, "another"+os.extsep+"name", self.data)

   144         zipfp.close()

   145 

   146     def testLargeFileException(self):

   147         for f in (TESTFN2, TemporaryFile(), StringIO()):

   148             self.largeFileExceptionTest(f, zipfile.ZIP_STORED)

   149             self.largeFileExceptionTest2(f, zipfile.ZIP_STORED)

   150 

   151     def zipTest(self, f, compression):

   152         # Create the ZIP archive

   153         zipfp = zipfile.ZipFile(f, "w", compression, allowZip64=True)

   154         zipfp.write(TESTFN, "another"+os.extsep+"name")

   155         zipfp.write(TESTFN, TESTFN)

   156         zipfp.writestr("strfile", self.data)

   157         zipfp.close()

   158 

   159         # Read the ZIP archive

   160         zipfp = zipfile.ZipFile(f, "r", compression)

   161         self.assertEqual(zipfp.read(TESTFN), self.data)

   162         self.assertEqual(zipfp.read("another"+os.extsep+"name"), self.data)

   163         self.assertEqual(zipfp.read("strfile"), self.data)

   164 

   165         # Print the ZIP directory

   166         fp = StringIO()

   167         stdout = sys.stdout

   168         try:

   169             sys.stdout = fp

   170 

   171             zipfp.printdir()

   172         finally:

   173             sys.stdout = stdout

   174 

   175         directory = fp.getvalue()

   176         lines = directory.splitlines()

   177         self.assertEquals(len(lines), 4) # Number of files + header

   178 

   179         self.assert_('File Name' in lines[0])

   180         self.assert_('Modified' in lines[0])

   181         self.assert_('Size' in lines[0])

   182 

   183         fn, date, time, size = lines[1].split()

   184         self.assertEquals(fn, 'another.name')

   185         # XXX: timestamp is not tested

   186         self.assertEquals(size, str(len(self.data)))

   187 

   188         # Check the namelist

   189         names = zipfp.namelist()

   190         self.assertEquals(len(names), 3)

   191         self.assert_(TESTFN in names)

   192         self.assert_("another"+os.extsep+"name" in names)

   193         self.assert_("strfile" in names)

   194 

   195         # Check infolist

   196         infos = zipfp.infolist()

   197         names = [ i.filename for i in infos ]

   198         self.assertEquals(len(names), 3)

   199         self.assert_(TESTFN in names)

   200         self.assert_("another"+os.extsep+"name" in names)

   201         self.assert_("strfile" in names)

   202         for i in infos:

   203             self.assertEquals(i.file_size, len(self.data))

   204 

   205         # check getinfo

   206         for nm in (TESTFN, "another"+os.extsep+"name", "strfile"):

   207             info = zipfp.getinfo(nm)

   208             self.assertEquals(info.filename, nm)

   209             self.assertEquals(info.file_size, len(self.data))

   210 

   211         # Check that testzip doesn't raise an exception

   212         zipfp.testzip()

   213 

   214 

   215         zipfp.close()

   216 

   217     def testStored(self):

   218         for f in (TESTFN2, TemporaryFile(), StringIO()):

   219             self.zipTest(f, zipfile.ZIP_STORED)

   220 

   221 

   222     if zlib:

   223         def testDeflated(self):

   224             for f in (TESTFN2, TemporaryFile(), StringIO()):

   225                 self.zipTest(f, zipfile.ZIP_DEFLATED)

   226 

   227     def testAbsoluteArcnames(self):

   228         zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED, allowZip64=True)

   229         zipfp.write(TESTFN, "/absolute")

   230         zipfp.close()

   231 

   232         zipfp = zipfile.ZipFile(TESTFN2, "r", zipfile.ZIP_STORED)

   233         self.assertEqual(zipfp.namelist(), ["absolute"])

   234         zipfp.close()

   235 

   236 

   237     def tearDown(self):

   238         zipfile.ZIP64_LIMIT = self._limit

   239         os.remove(TESTFN)

   240         os.remove(TESTFN2)

   241 

   242 class PyZipFileTests(unittest.TestCase):

   243     def testWritePyfile(self):

   244         zipfp  = zipfile.PyZipFile(TemporaryFile(), "w")

   245         fn = __file__

   246         if fn.endswith('.pyc') or fn.endswith('.pyo'):

   247             fn = fn[:-1]

   248 

   249         zipfp.writepy(fn)

   250 

   251         bn = os.path.basename(fn)

   252         self.assert_(bn not in zipfp.namelist())

   253         self.assert_(bn + 'o' in zipfp.namelist() or bn + 'c' in zipfp.namelist())

   254         zipfp.close()

   255 

   256 

   257         zipfp  = zipfile.PyZipFile(TemporaryFile(), "w")

   258         fn = __file__

   259         if fn.endswith('.pyc') or fn.endswith('.pyo'):

   260             fn = fn[:-1]

   261 

   262         zipfp.writepy(fn, "testpackage")

   263 

   264         bn = "%s/%s"%("testpackage", os.path.basename(fn))

   265         self.assert_(bn not in zipfp.namelist())

   266         self.assert_(bn + 'o' in zipfp.namelist() or bn + 'c' in zipfp.namelist())

   267         zipfp.close()

   268 

   269     def testWritePythonPackage(self):

   270         import email

   271         packagedir = os.path.dirname(email.__file__)

   272 

   273         zipfp  = zipfile.PyZipFile(TemporaryFile(), "w")

   274         zipfp.writepy(packagedir)

   275 

   276         # Check for a couple of modules at different levels of the hieararchy

   277         names = zipfp.namelist()

   278         self.assert_('email/__init__.pyo' in names or 'email/__init__.pyc' in names)

   279         self.assert_('email/mime/text.pyo' in names or 'email/mime/text.pyc' in names)

   280 

   281     def testWritePythonDirectory(self):

   282         os.mkdir(TESTFN2)

   283         try:

   284             fp = open(os.path.join(TESTFN2, "mod1.py"), "w")

   285             fp.write("print 42\n")

   286             fp.close()

   287 

   288             fp = open(os.path.join(TESTFN2, "mod2.py"), "w")

   289             fp.write("print 42 * 42\n")

   290             fp.close()

   291 

   292             fp = open(os.path.join(TESTFN2, "mod2.txt"), "w")

   293             fp.write("bla bla bla\n")

   294             fp.close()

   295 

   296             zipfp  = zipfile.PyZipFile(TemporaryFile(), "w")

   297             zipfp.writepy(TESTFN2)

   298 

   299             names = zipfp.namelist()

   300             self.assert_('mod1.pyc' in names or 'mod1.pyo' in names)

   301             self.assert_('mod2.pyc' in names or 'mod2.pyo' in names)

   302             self.assert_('mod2.txt' not in names)

   303 

   304         finally:

   305             shutil.rmtree(TESTFN2)

   306 

   307 

   308 

   309 class OtherTests(unittest.TestCase):

   310     def testCloseErroneousFile(self):

   311         # This test checks that the ZipFile constructor closes the file object

   312         # it opens if there's an error in the file.  If it doesn't, the traceback

   313         # holds a reference to the ZipFile object and, indirectly, the file object.

   314         # On Windows, this causes the os.unlink() call to fail because the

   315         # underlying file is still open.  This is SF bug #412214.

   316         #

   317         fp = open(TESTFN, "w")

   318         fp.write("this is not a legal zip file\n")

   319         fp.close()

   320         try:

   321             zf = zipfile.ZipFile(TESTFN)

   322         except zipfile.BadZipfile:

   323             os.unlink(TESTFN)

   324 

   325     def testNonExistentFileRaisesIOError(self):

   326         # make sure we don't raise an AttributeError when a partially-constructed

   327         # ZipFile instance is finalized; this tests for regression on SF tracker

   328         # bug #403871.

   329 

   330         # The bug we're testing for caused an AttributeError to be raised

   331         # when a ZipFile instance was created for a file that did not

   332         # exist; the .fp member was not initialized but was needed by the

   333         # __del__() method.  Since the AttributeError is in the __del__(),

   334         # it is ignored, but the user should be sufficiently annoyed by

   335         # the message on the output that regression will be noticed

   336         # quickly.

   337         self.assertRaises(IOError, zipfile.ZipFile, TESTFN)

   338 

   339     def testClosedZipRaisesRuntimeError(self):

   340         # Verify that testzip() doesn't swallow inappropriate exceptions.

   341         data = StringIO()

   342         zipf = zipfile.ZipFile(data, mode="w")

   343         zipf.writestr("foo.txt", "O, for a Muse of Fire!")

   344         zipf.close()

   345 

   346         # This is correct; calling .read on a closed ZipFile should throw

   347         # a RuntimeError, and so should calling .testzip.  An earlier

   348         # version of .testzip would swallow this exception (and any other)

   349         # and report that the first file in the archive was corrupt.

   350         self.assertRaises(RuntimeError, zipf.testzip)

   351 

   352 def test_main():

   353     run_unittest(TestsWithSourceFile, TestZip64InSmallFiles, OtherTests, PyZipFileTests)

   354     #run_unittest(TestZip64InSmallFiles)

   355 

   356 if __name__ == "__main__":

   357     test_main()