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