Lib/test/test_zipimport.py

   1 import sys
   2 import os
   3 import marshal
   4 import imp
   5 import struct
   6 import time
   7
   8 import zlib # implied prerequisite
   9 from zipfile import ZipFile, ZipInfo, ZIP_STORED, ZIP_DEFLATED
  10 from test import test_support
  11 from test.test_importhooks import ImportHooksBaseTestCase, test_src, test_co
  12
  13 import zipimport
  14
  15
  16 def make_pyc(co, mtime):
  17     data = marshal.dumps(co)
  18     if type(mtime) is type(0.0):
  19         # Mac mtimes need a bit of special casing
  20         if mtime < 0x7fffffff:
  21             mtime = int(mtime)
  22         else:
  23             mtime = int(-0x100000000L + long(mtime))
  24     pyc = imp.get_magic() + struct.pack("<i", int(mtime)) + data
  25     return pyc
  26
  27 NOW = time.time()
  28 test_pyc = make_pyc(test_co, NOW)
  29
  30
  31 if __debug__:
  32     pyc_ext = ".pyc"
  33 else:
  34     pyc_ext = ".pyo"
  35
  36
  37 TESTMOD = "ziptestmodule"
  38 TESTPACK = "ziptestpackage"
  39 TESTPACK2 = "ziptestpackage2"
  40 TEMP_ZIP = os.path.abspath("junk95142.zip")
  41
  42 class UncompressedZipImportTestCase(ImportHooksBaseTestCase):
  43
  44     compression = ZIP_STORED
  45
  46     def setUp(self):
  47         # We're reusing the zip archive path, so we must clear the
  48         # cached directory info.
  49         zipimport._zip_directory_cache.clear()
  50         ImportHooksBaseTestCase.setUp(self)
  51
  52     def doTest(self, expected_ext, files, *modules):
  53         z = ZipFile(TEMP_ZIP, "w")
  54         try:
  55             for name, (mtime, data) in files.items():
  56                 zinfo = ZipInfo(name, time.localtime(mtime))
  57                 zinfo.compress_type = self.compression
  58                 z.writestr(zinfo, data)
  59             z.close()
  60             sys.path.insert(0, TEMP_ZIP)
  61
  62             mod = __import__(".".join(modules), globals(), locals(),
  63                              ["__dummy__"])
  64             if expected_ext:
  65                 file = mod.get_file()
  66                 self.assertEquals(file, os.path.join(TEMP_ZIP,
  67                                   *modules) + expected_ext)
  68         finally:
  69             z.close()
  70             os.remove(TEMP_ZIP)
  71
  72     def testAFakeZlib(self):
  73         #
  74         # This could cause a stack overflow before: importing zlib.py
  75         # from a compressed archive would cause zlib to be imported
  76         # which would find zlib.py in the archive, which would... etc.
  77         #
  78         # This test *must* be executed first: it must be the first one
  79         # to trigger zipimport to import zlib (zipimport caches the
  80         # zlib.decompress function object, after which the problem being
  81         # tested here wouldn't be a problem anymore...
  82         # (Hence the 'A' in the test method name: to make it the first
  83         # item in a list sorted by name, like unittest.makeSuite() does.)
  84         #
  85         if "zlib" in sys.modules:
  86             del sys.modules["zlib"]
  87         files = {"zlib.py": (NOW, test_src)}
  88         try:
  89             self.doTest(".py", files, "zlib")
  90         except ImportError:
  91             if self.compression != ZIP_DEFLATED:
  92                 self.fail("expected test to not raise ImportError")
  93         else:
  94             if self.compression != ZIP_STORED:
  95                 self.fail("expected test to raise ImportError")
  96
  97     def testPy(self):
  98         files = {TESTMOD + ".py": (NOW, test_src)}
  99         self.doTest(".py", files, TESTMOD)
 100
 101     def testPyc(self):
 102         files = {TESTMOD + pyc_ext: (NOW, test_pyc)}
 103         self.doTest(pyc_ext, files, TESTMOD)
 104
 105     def testBoth(self):
 106         files = {TESTMOD + ".py": (NOW, test_src),
 107                  TESTMOD + pyc_ext: (NOW, test_pyc)}
 108         self.doTest(pyc_ext, files, TESTMOD)
 109
 110     def testEmptyPy(self):
 111         files = {TESTMOD + ".py": (NOW, "")}
 112         self.doTest(None, files, TESTMOD)
 113
 114     def testBadMagic(self):
 115         # make pyc magic word invalid, forcing loading from .py
 116         m0 = ord(test_pyc[0])
 117         m0 ^= 0x04  # flip an arbitrary bit
 118         badmagic_pyc = chr(m0) + test_pyc[1:]
 119         files = {TESTMOD + ".py": (NOW, test_src),
 120                  TESTMOD + pyc_ext: (NOW, badmagic_pyc)}
 121         self.doTest(".py", files, TESTMOD)
 122
 123     def testBadMagic2(self):
 124         # make pyc magic word invalid, causing an ImportError
 125         m0 = ord(test_pyc[0])
 126         m0 ^= 0x04  # flip an arbitrary bit
 127         badmagic_pyc = chr(m0) + test_pyc[1:]
 128         files = {TESTMOD + pyc_ext: (NOW, badmagic_pyc)}
 129         try:
 130             self.doTest(".py", files, TESTMOD)
 131         except ImportError:
 132             pass
 133         else:
 134             self.fail("expected ImportError; import from bad pyc")
 135
 136     def testBadMTime(self):
 137         t3 = ord(test_pyc[7])
 138         t3 ^= 0x02  # flip the second bit -- not the first as that one
 139                     # isn't stored in the .py's mtime in the zip archive.
 140         badtime_pyc = test_pyc[:7] + chr(t3) + test_pyc[8:]
 141         files = {TESTMOD + ".py": (NOW, test_src),
 142                  TESTMOD + pyc_ext: (NOW, badtime_pyc)}
 143         self.doTest(".py", files, TESTMOD)
 144
 145     def testPackage(self):
 146         packdir = TESTPACK + os.sep
 147         files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
 148                  packdir + TESTMOD + pyc_ext: (NOW, test_pyc)}
 149         self.doTest(pyc_ext, files, TESTPACK, TESTMOD)
 150
 151     def testDeepPackage(self):
 152         packdir = TESTPACK + os.sep
 153         packdir2 = packdir + TESTPACK2 + os.sep
 154         files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
 155                  packdir2 + "__init__" + pyc_ext: (NOW, test_pyc),
 156                  packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc)}
 157         self.doTest(pyc_ext, files, TESTPACK, TESTPACK2, TESTMOD)
 158
 159     def testGetData(self):
 160         z = ZipFile(TEMP_ZIP, "w")
 161         z.compression = self.compression
 162         try:
 163             name = "testdata.dat"
 164             data = "".join([chr(x) for x in range(256)]) * 500
 165             z.writestr(name, data)
 166             z.close()
 167             zi = zipimport.zipimporter(TEMP_ZIP)
 168             self.assertEquals(data, zi.get_data(name))
 169         finally:
 170             z.close()
 171             os.remove(TEMP_ZIP)
 172
 173     def testImporterAttr(self):
 174         src = """if 1:  # indent hack
 175         def get_file():
 176             return __file__
 177         if __loader__.get_data("some.data") != "some data":
 178             raise AssertionError, "bad data"\n"""
 179         pyc = make_pyc(compile(src, "<???>", "exec"), NOW)
 180         files = {TESTMOD + pyc_ext: (NOW, pyc),
 181                  "some.data": (NOW, "some data")}
 182         self.doTest(pyc_ext, files, TESTMOD)
 183
 184
 185 class CompressedZipImportTestCase(UncompressedZipImportTestCase):
 186     compression = ZIP_DEFLATED
 187
 188
 189 def test_main():
 190     test_support.run_unittest(UncompressedZipImportTestCase)
 191     test_support.run_unittest(CompressedZipImportTestCase)
 192
 193 if __name__ == "__main__":
 194     test_main()