Lib/test/test_zipimport.py

   1 import sys
   2 import os
   3 import marshal
   4 import imp
   5 import struct
   6 import time
   7 import unittest
   8
   9 from test import test_support
  10 from test.test_importhooks import ImportHooksBaseTestCase, test_src, test_co
  11
  12 # some tests can be ran even without zlib
  13 try:
  14     import zlib
  15 except ImportError:
  16     zlib = None
  17
  18 from zipfile import ZipFile, ZipInfo, ZIP_STORED, ZIP_DEFLATED
  19
  20 import zipimport
  21 import linecache
  22 import doctest
  23 import inspect
  24 import StringIO
  25 from traceback import extract_tb, extract_stack, print_tb
  26 raise_src = 'def do_raise(): raise TypeError\n'
  27
  28 def make_pyc(co, mtime):
  29     data = marshal.dumps(co)
  30     if type(mtime) is type(0.0):
  31         # Mac mtimes need a bit of special casing
  32         if mtime < 0x7fffffff:
  33             mtime = int(mtime)
  34         else:
  35             mtime = int(-0x100000000L + long(mtime))
  36     pyc = imp.get_magic() + struct.pack("<i", int(mtime)) + data
  37     return pyc
  38
  39 def module_path_to_dotted_name(path):
  40     return path.replace(os.sep, '.')
  41
  42 NOW = time.time()
  43 test_pyc = make_pyc(test_co, NOW)
  44
  45
  46 if __debug__:
  47     pyc_ext = ".pyc"
  48 else:
  49     pyc_ext = ".pyo"
  50
  51
  52 TESTMOD = "ziptestmodule"
  53 TESTPACK = "ziptestpackage"
  54 TESTPACK2 = "ziptestpackage2"
  55 TEMP_ZIP = os.path.abspath("junk95142" + os.extsep + "zip")
  56
  57
  58 class UncompressedZipImportTestCase(ImportHooksBaseTestCase):
  59
  60     compression = ZIP_STORED
  61
  62     def setUp(self):
  63         # We're reusing the zip archive path, so we must clear the
  64         # cached directory info and linecache
  65         linecache.clearcache()
  66         zipimport._zip_directory_cache.clear()
  67         ImportHooksBaseTestCase.setUp(self)
  68
  69     def doTest(self, expected_ext, files, *modules, **kw):
  70         z = ZipFile(TEMP_ZIP, "w")
  71         try:
  72             for name, (mtime, data) in files.items():
  73                 zinfo = ZipInfo(name, time.localtime(mtime))
  74                 zinfo.compress_type = self.compression
  75                 z.writestr(zinfo, data)
  76             z.close()
  77
  78             stuff = kw.get("stuff", None)
  79             if stuff is not None:
  80                 # Prepend 'stuff' to the start of the zipfile
  81                 f = open(TEMP_ZIP, "rb")
  82                 data = f.read()
  83                 f.close()
  84
  85                 f = open(TEMP_ZIP, "wb")
  86                 f.write(stuff)
  87                 f.write(data)
  88                 f.close()
  89
  90             sys.path.insert(0, TEMP_ZIP)
  91
  92             mod = __import__(".".join(modules), globals(), locals(),
  93                              ["__dummy__"])
  94
  95             call = kw.get('call')
  96             if call is not None:
  97                 call(mod)
  98
  99             if expected_ext:
 100                 file = mod.get_file()
 101                 self.assertEqual(file, os.path.join(TEMP_ZIP,
 102                                  *modules) + expected_ext)
 103         finally:
 104             z.close()
 105             os.remove(TEMP_ZIP)
 106
 107     def testAFakeZlib(self):
 108         #
 109         # This could cause a stack overflow before: importing zlib.py
 110         # from a compressed archive would cause zlib to be imported
 111         # which would find zlib.py in the archive, which would... etc.
 112         #
 113         # This test *must* be executed first: it must be the first one
 114         # to trigger zipimport to import zlib (zipimport caches the
 115         # zlib.decompress function object, after which the problem being
 116         # tested here wouldn't be a problem anymore...
 117         # (Hence the 'A' in the test method name: to make it the first
 118         # item in a list sorted by name, like unittest.makeSuite() does.)
 119         #
 120         # This test fails on platforms on which the zlib module is
 121         # statically linked, but the problem it tests for can't
 122         # occur in that case (builtin modules are always found first),
 123         # so we'll simply skip it then. Bug #765456.
 124         #
 125         if "zlib" in sys.builtin_module_names:
 126             return
 127         if "zlib" in sys.modules:
 128             del sys.modules["zlib"]
 129         files = {"zlib.py": (NOW, test_src)}
 130         try:
 131             self.doTest(".py", files, "zlib")
 132         except ImportError:
 133             if self.compression != ZIP_DEFLATED:
 134                 self.fail("expected test to not raise ImportError")
 135         else:
 136             if self.compression != ZIP_STORED:
 137                 self.fail("expected test to raise ImportError")
 138
 139     def testPy(self):
 140         files = {TESTMOD + ".py": (NOW, test_src)}
 141         self.doTest(".py", files, TESTMOD)
 142
 143     def testPyc(self):
 144         files = {TESTMOD + pyc_ext: (NOW, test_pyc)}
 145         self.doTest(pyc_ext, files, TESTMOD)
 146
 147     def testBoth(self):
 148         files = {TESTMOD + ".py": (NOW, test_src),
 149                  TESTMOD + pyc_ext: (NOW, test_pyc)}
 150         self.doTest(pyc_ext, files, TESTMOD)
 151
 152     def testEmptyPy(self):
 153         files = {TESTMOD + ".py": (NOW, "")}
 154         self.doTest(None, files, TESTMOD)
 155
 156     def testBadMagic(self):
 157         # make pyc magic word invalid, forcing loading from .py
 158         m0 = ord(test_pyc[0])
 159         m0 ^= 0x04  # flip an arbitrary bit
 160         badmagic_pyc = chr(m0) + test_pyc[1:]
 161         files = {TESTMOD + ".py": (NOW, test_src),
 162                  TESTMOD + pyc_ext: (NOW, badmagic_pyc)}
 163         self.doTest(".py", files, TESTMOD)
 164
 165     def testBadMagic2(self):
 166         # make pyc magic word invalid, causing an ImportError
 167         m0 = ord(test_pyc[0])
 168         m0 ^= 0x04  # flip an arbitrary bit
 169         badmagic_pyc = chr(m0) + test_pyc[1:]
 170         files = {TESTMOD + pyc_ext: (NOW, badmagic_pyc)}
 171         try:
 172             self.doTest(".py", files, TESTMOD)
 173         except ImportError:
 174             pass
 175         else:
 176             self.fail("expected ImportError; import from bad pyc")
 177
 178     def testBadMTime(self):
 179         t3 = ord(test_pyc[7])
 180         t3 ^= 0x02  # flip the second bit -- not the first as that one
 181                     # isn't stored in the .py's mtime in the zip archive.
 182         badtime_pyc = test_pyc[:7] + chr(t3) + test_pyc[8:]
 183         files = {TESTMOD + ".py": (NOW, test_src),
 184                  TESTMOD + pyc_ext: (NOW, badtime_pyc)}
 185         self.doTest(".py", files, TESTMOD)
 186
 187     def testPackage(self):
 188         packdir = TESTPACK + os.sep
 189         files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
 190                  packdir + TESTMOD + pyc_ext: (NOW, test_pyc)}
 191         self.doTest(pyc_ext, files, TESTPACK, TESTMOD)
 192
 193     def testDeepPackage(self):
 194         packdir = TESTPACK + os.sep
 195         packdir2 = packdir + TESTPACK2 + os.sep
 196         files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
 197                  packdir2 + "__init__" + pyc_ext: (NOW, test_pyc),
 198                  packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc)}
 199         self.doTest(pyc_ext, files, TESTPACK, TESTPACK2, TESTMOD)
 200
 201     def testZipImporterMethods(self):
 202         packdir = TESTPACK + os.sep
 203         packdir2 = packdir + TESTPACK2 + os.sep
 204         files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
 205                  packdir2 + "__init__" + pyc_ext: (NOW, test_pyc),
 206                  packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc)}
 207
 208         z = ZipFile(TEMP_ZIP, "w")
 209         try:
 210             for name, (mtime, data) in files.items():
 211                 zinfo = ZipInfo(name, time.localtime(mtime))
 212                 zinfo.compress_type = self.compression
 213                 z.writestr(zinfo, data)
 214             z.close()
 215
 216             zi = zipimport.zipimporter(TEMP_ZIP)
 217             self.assertEqual(zi.archive, TEMP_ZIP)
 218             self.assertEqual(zi.is_package(TESTPACK), True)
 219             mod = zi.load_module(TESTPACK)
 220             self.assertEqual(zi.get_filename(TESTPACK), mod.__file__)
 221
 222             self.assertEqual(zi.is_package(packdir + '__init__'), False)
 223             self.assertEqual(zi.is_package(packdir + TESTPACK2), True)
 224             self.assertEqual(zi.is_package(packdir2 + TESTMOD), False)
 225
 226             mod_path = packdir2 + TESTMOD
 227             mod_name = module_path_to_dotted_name(mod_path)
 228             __import__(mod_name)
 229             mod = sys.modules[mod_name]
 230             self.assertEqual(zi.get_source(TESTPACK), None)
 231             self.assertEqual(zi.get_source(mod_path), None)
 232             self.assertEqual(zi.get_filename(mod_path), mod.__file__)
 233             # To pass in the module name instead of the path, we must use the right importer
 234             loader = mod.__loader__
 235             self.assertEqual(loader.get_source(mod_name), None)
 236             self.assertEqual(loader.get_filename(mod_name), mod.__file__)
 237
 238             # test prefix and archivepath members
 239             zi2 = zipimport.zipimporter(TEMP_ZIP + os.sep + TESTPACK)
 240             self.assertEqual(zi2.archive, TEMP_ZIP)
 241             self.assertEqual(zi2.prefix, TESTPACK + os.sep)
 242         finally:
 243             z.close()
 244             os.remove(TEMP_ZIP)
 245
 246     def testZipImporterMethodsInSubDirectory(self):
 247         packdir = TESTPACK + os.sep
 248         packdir2 = packdir + TESTPACK2 + os.sep
 249         files = {packdir2 + "__init__" + pyc_ext: (NOW, test_pyc),
 250                  packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc)}
 251
 252         z = ZipFile(TEMP_ZIP, "w")
 253         try:
 254             for name, (mtime, data) in files.items():
 255                 zinfo = ZipInfo(name, time.localtime(mtime))
 256                 zinfo.compress_type = self.compression
 257                 z.writestr(zinfo, data)
 258             z.close()
 259
 260             zi = zipimport.zipimporter(TEMP_ZIP + os.sep + packdir)
 261             self.assertEqual(zi.archive, TEMP_ZIP)
 262             self.assertEqual(zi.prefix, packdir)
 263             self.assertEqual(zi.is_package(TESTPACK2), True)
 264             mod = zi.load_module(TESTPACK2)
 265             self.assertEqual(zi.get_filename(TESTPACK2), mod.__file__)
 266
 267             self.assertEqual(zi.is_package(TESTPACK2 + os.sep + '__init__'), False)
 268             self.assertEqual(zi.is_package(TESTPACK2 + os.sep + TESTMOD), False)
 269
 270             mod_path = TESTPACK2 + os.sep + TESTMOD
 271             mod_name = module_path_to_dotted_name(mod_path)
 272             __import__(mod_name)
 273             mod = sys.modules[mod_name]
 274             self.assertEqual(zi.get_source(TESTPACK2), None)
 275             self.assertEqual(zi.get_source(mod_path), None)
 276             self.assertEqual(zi.get_filename(mod_path), mod.__file__)
 277             # To pass in the module name instead of the path, we must use the right importer
 278             loader = mod.__loader__
 279             self.assertEqual(loader.get_source(mod_name), None)
 280             self.assertEqual(loader.get_filename(mod_name), mod.__file__)
 281         finally:
 282             z.close()
 283             os.remove(TEMP_ZIP)
 284
 285     def testGetData(self):
 286         z = ZipFile(TEMP_ZIP, "w")
 287         z.compression = self.compression
 288         try:
 289             name = "testdata.dat"
 290             data = "".join([chr(x) for x in range(256)]) * 500
 291             z.writestr(name, data)
 292             z.close()
 293             zi = zipimport.zipimporter(TEMP_ZIP)
 294             self.assertEqual(data, zi.get_data(name))
 295             self.assertIn('zipimporter object', repr(zi))
 296         finally:
 297             z.close()
 298             os.remove(TEMP_ZIP)
 299
 300     def testImporterAttr(self):
 301         src = """if 1:  # indent hack
 302         def get_file():
 303             return __file__
 304         if __loader__.get_data("some.data") != "some data":
 305             raise AssertionError, "bad data"\n"""
 306         pyc = make_pyc(compile(src, "<???>", "exec"), NOW)
 307         files = {TESTMOD + pyc_ext: (NOW, pyc),
 308                  "some.data": (NOW, "some data")}
 309         self.doTest(pyc_ext, files, TESTMOD)
 310
 311     def testImport_WithStuff(self):
 312         # try importing from a zipfile which contains additional
 313         # stuff at the beginning of the file
 314         files = {TESTMOD + ".py": (NOW, test_src)}
 315         self.doTest(".py", files, TESTMOD,
 316                     stuff="Some Stuff"*31)
 317
 318     def assertModuleSource(self, module):
 319         self.assertEqual(inspect.getsource(module), test_src)
 320
 321     def testGetSource(self):
 322         files = {TESTMOD + ".py": (NOW, test_src)}
 323         self.doTest(".py", files, TESTMOD, call=self.assertModuleSource)
 324
 325     def testGetCompiledSource(self):
 326         pyc = make_pyc(compile(test_src, "<???>", "exec"), NOW)
 327         files = {TESTMOD + ".py": (NOW, test_src),
 328                  TESTMOD + pyc_ext: (NOW, pyc)}
 329         self.doTest(pyc_ext, files, TESTMOD, call=self.assertModuleSource)
 330
 331     def runDoctest(self, callback):
 332         files = {TESTMOD + ".py": (NOW, test_src),
 333                  "xyz.txt": (NOW, ">>> log.append(True)\n")}
 334         self.doTest(".py", files, TESTMOD, call=callback)
 335
 336     def doDoctestFile(self, module):
 337         log = []
 338         old_master, doctest.master = doctest.master, None
 339         try:
 340             doctest.testfile(
 341                 'xyz.txt', package=module, module_relative=True,
 342                 globs=locals()
 343             )
 344         finally:
 345             doctest.master = old_master
 346         self.assertEqual(log,[True])
 347
 348     def testDoctestFile(self):
 349         self.runDoctest(self.doDoctestFile)
 350
 351     def doDoctestSuite(self, module):
 352         log = []
 353         doctest.DocFileTest(
 354             'xyz.txt', package=module, module_relative=True,
 355             globs=locals()
 356         ).run()
 357         self.assertEqual(log,[True])
 358
 359     def testDoctestSuite(self):
 360         self.runDoctest(self.doDoctestSuite)
 361
 362     def doTraceback(self, module):
 363         try:
 364             module.do_raise()
 365         except:
 366             tb = sys.exc_info()[2].tb_next
 367
 368             f,lno,n,line = extract_tb(tb, 1)[0]
 369             self.assertEqual(line, raise_src.strip())
 370
 371             f,lno,n,line = extract_stack(tb.tb_frame, 1)[0]
 372             self.assertEqual(line, raise_src.strip())
 373
 374             s = StringIO.StringIO()
 375             print_tb(tb, 1, s)
 376             self.assertTrue(s.getvalue().endswith(raise_src))
 377         else:
 378             raise AssertionError("This ought to be impossible")
 379
 380     def testTraceback(self):
 381         files = {TESTMOD + ".py": (NOW, raise_src)}
 382         self.doTest(None, files, TESTMOD, call=self.doTraceback)
 383
 384
 385 @unittest.skipUnless(zlib, "requires zlib")
 386 class CompressedZipImportTestCase(UncompressedZipImportTestCase):
 387     compression = ZIP_DEFLATED
 388
 389
 390 class BadFileZipImportTestCase(unittest.TestCase):
 391     def assertZipFailure(self, filename):
 392         self.assertRaises(zipimport.ZipImportError,
 393                           zipimport.zipimporter, filename)
 394
 395     def testNoFile(self):
 396         self.assertZipFailure('AdfjdkFJKDFJjdklfjs')
 397
 398     def testEmptyFilename(self):
 399         self.assertZipFailure('')
 400
 401     def testBadArgs(self):
 402         self.assertRaises(TypeError, zipimport.zipimporter, None)
 403         self.assertRaises(TypeError, zipimport.zipimporter, TESTMOD, kwd=None)
 404
 405     def testFilenameTooLong(self):
 406         self.assertZipFailure('A' * 33000)
 407
 408     def testEmptyFile(self):
 409         test_support.unlink(TESTMOD)
 410         open(TESTMOD, 'w+').close()
 411         self.assertZipFailure(TESTMOD)
 412
 413     def testFileUnreadable(self):
 414         test_support.unlink(TESTMOD)
 415         fd = os.open(TESTMOD, os.O_CREAT, 000)
 416         try:
 417             os.close(fd)
 418             self.assertZipFailure(TESTMOD)
 419         finally:
 420             # If we leave "the read-only bit" set on Windows, nothing can
 421             # delete TESTMOD, and later tests suffer bogus failures.
 422             os.chmod(TESTMOD, 0666)
 423             test_support.unlink(TESTMOD)
 424
 425     def testNotZipFile(self):
 426         test_support.unlink(TESTMOD)
 427         fp = open(TESTMOD, 'w+')
 428         fp.write('a' * 22)
 429         fp.close()
 430         self.assertZipFailure(TESTMOD)
 431
 432     # XXX: disabled until this works on Big-endian machines
 433     def _testBogusZipFile(self):
 434         test_support.unlink(TESTMOD)
 435         fp = open(TESTMOD, 'w+')
 436         fp.write(struct.pack('=I', 0x06054B50))
 437         fp.write('a' * 18)
 438         fp.close()
 439         z = zipimport.zipimporter(TESTMOD)
 440
 441         try:
 442             self.assertRaises(TypeError, z.find_module, None)
 443             self.assertRaises(TypeError, z.load_module, None)
 444             self.assertRaises(TypeError, z.is_package, None)
 445             self.assertRaises(TypeError, z.get_code, None)
 446             self.assertRaises(TypeError, z.get_data, None)
 447             self.assertRaises(TypeError, z.get_source, None)
 448
 449             error = zipimport.ZipImportError
 450             self.assertEqual(z.find_module('abc'), None)
 451
 452             self.assertRaises(error, z.load_module, 'abc')
 453             self.assertRaises(error, z.get_code, 'abc')
 454             self.assertRaises(IOError, z.get_data, 'abc')
 455             self.assertRaises(error, z.get_source, 'abc')
 456             self.assertRaises(error, z.is_package, 'abc')
 457         finally:
 458             zipimport._zip_directory_cache.clear()
 459
 460
 461 def test_main():
 462     try:
 463         test_support.run_unittest(
 464               UncompressedZipImportTestCase,
 465               CompressedZipImportTestCase,
 466               BadFileZipImportTestCase,
 467             )
 468     finally:
 469         test_support.unlink(TESTMOD)
 470
 471 if __name__ == "__main__":
 472     test_main()