import sys import os import marshal import importlib.util import struct import time import unittest from test import support from zipfile import ZipFile, ZipInfo, ZIP_STORED, ZIP_DEFLATED import zipimport import linecache import doctest import inspect import io from traceback import extract_tb, extract_stack, print_tb test_src = """\ def get_name(): return __name__ def get_file(): return __file__ """ test_co = compile(test_src, "", "exec") raise_src = 'def do_raise(): raise TypeError\n' def make_pyc(co, mtime, size): data = marshal.dumps(co) if type(mtime) is type(0.0): # Mac mtimes need a bit of special casing if mtime < 0x7fffffff: mtime = int(mtime) else: mtime = int(-0x100000000 + int(mtime)) pyc = (importlib.util.MAGIC_NUMBER + struct.pack("", "exec"), NOW, len(src)) files = {TESTMOD + pyc_ext: (NOW, pyc), "some.data": (NOW, "some data")} self.doTest(pyc_ext, files, TESTMOD) def testImport_WithStuff(self): # try importing from a zipfile which contains additional # stuff at the beginning of the file files = {TESTMOD + ".py": (NOW, test_src)} self.doTest(".py", files, TESTMOD, stuff=b"Some Stuff"*31) def assertModuleSource(self, module): self.assertEqual(inspect.getsource(module), test_src) def testGetSource(self): files = {TESTMOD + ".py": (NOW, test_src)} self.doTest(".py", files, TESTMOD, call=self.assertModuleSource) def testGetCompiledSource(self): pyc = make_pyc(compile(test_src, "", "exec"), NOW, len(test_src)) files = {TESTMOD + ".py": (NOW, test_src), TESTMOD + pyc_ext: (NOW, pyc)} self.doTest(pyc_ext, files, TESTMOD, call=self.assertModuleSource) def runDoctest(self, callback): files = {TESTMOD + ".py": (NOW, test_src), "xyz.txt": (NOW, ">>> log.append(True)\n")} self.doTest(".py", files, TESTMOD, call=callback) def doDoctestFile(self, module): log = [] old_master, doctest.master = doctest.master, None try: doctest.testfile( 'xyz.txt', package=module, module_relative=True, globs=locals() ) finally: doctest.master = old_master self.assertEqual(log,[True]) def testDoctestFile(self): self.runDoctest(self.doDoctestFile) def doDoctestSuite(self, module): log = [] doctest.DocFileTest( 'xyz.txt', package=module, module_relative=True, globs=locals() ).run() self.assertEqual(log,[True]) def testDoctestSuite(self): self.runDoctest(self.doDoctestSuite) def doTraceback(self, module): try: module.do_raise() except: tb = sys.exc_info()[2].tb_next f,lno,n,line = extract_tb(tb, 1)[0] self.assertEqual(line, raise_src.strip()) f,lno,n,line = extract_stack(tb.tb_frame, 1)[0] self.assertEqual(line, raise_src.strip()) s = io.StringIO() print_tb(tb, 1, s) self.assertTrue(s.getvalue().endswith(raise_src)) else: raise AssertionError("This ought to be impossible") def testTraceback(self): files = {TESTMOD + ".py": (NOW, raise_src)} self.doTest(None, files, TESTMOD, call=self.doTraceback) @unittest.skipIf(support.TESTFN_UNENCODABLE is None, "need an unencodable filename") def testUnencodable(self): filename = support.TESTFN_UNENCODABLE + ".zip" z = ZipFile(filename, "w") zinfo = ZipInfo(TESTMOD + ".py", time.localtime(NOW)) zinfo.compress_type = self.compression z.writestr(zinfo, test_src) z.close() try: zipimport.zipimporter(filename) finally: os.remove(filename) @support.requires_zlib class CompressedZipImportTestCase(UncompressedZipImportTestCase): compression = ZIP_DEFLATED class ZipFileModifiedAfterImportTestCase(ImportHooksBaseTestCase): def setUp(self): zipimport._zip_directory_cache.clear() zipimport._zip_stat_cache.clear() ImportHooksBaseTestCase.setUp(self) def tearDown(self): ImportHooksBaseTestCase.tearDown(self) if os.path.exists(TEMP_ZIP): os.remove(TEMP_ZIP) def testZipFileChangesAfterFirstImport(self): """Alter the zip file after caching its index and try an import.""" packdir = TESTPACK + os.sep files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc), packdir + TESTMOD + ".py": (NOW, "test_value = 38\n"), "ziptest_a.py": (NOW, "test_value = 23\n"), "ziptest_b.py": (NOW, "test_value = 42\n"), "ziptest_c.py": (NOW, "test_value = 1337\n")} zipfile_path = TEMP_ZIP _write_zip_package(zipfile_path, files) self.assertTrue(os.path.exists(zipfile_path)) sys.path.insert(0, zipfile_path) # Import something out of the zipfile and confirm it is correct. testmod = __import__(TESTPACK + "." + TESTMOD, globals(), locals(), ["__dummy__"]) self.assertEqual(testmod.test_value, 38) # Import something else out of the zipfile and confirm it is correct. ziptest_b = __import__("ziptest_b", globals(), locals(), ["test_value"]) self.assertEqual(ziptest_b.test_value, 42) # Truncate and fill the zip file with non-zip garbage. with open(zipfile_path, "rb") as orig_zip_file: orig_zip_file_contents = orig_zip_file.read() with open(zipfile_path, "wb") as byebye_valid_zip_file: byebye_valid_zip_file.write(b"Tear down this wall!\n"*1987) # Now that the zipfile has been replaced, import something else from it # which should fail as the file contents are now garbage. with self.assertRaises(ImportError): ziptest_a = __import__("ziptest_a", globals(), locals(), ["test_value"]) self.assertEqual(ziptest_a.test_value, 23) # Now lets make it a valid zipfile that has some garbage at the start. # This alters all of the offsets within the file with open(zipfile_path, "wb") as new_zip_file: new_zip_file.write(b"X"*1991) # The year Python was created. new_zip_file.write(orig_zip_file_contents) # Now that the zip file has been "restored" to a valid but different # zipfile the zipimporter should *successfully* re-read the new zip # file's end of file central index and be able to import from it again. ziptest_c = __import__("ziptest_c", globals(), locals(), ["test_value"]) self.assertEqual(ziptest_c.test_value, 1337) class BadFileZipImportTestCase(unittest.TestCase): def assertZipFailure(self, filename): with self.assertRaises(zipimport.ZipImportError): zipimport.zipimporter(filename) def testNoFile(self): self.assertZipFailure('AdfjdkFJKDFJjdklfjs') def testEmptyFilename(self): self.assertZipFailure('') def testBadArgs(self): self.assertRaises(TypeError, zipimport.zipimporter, None) self.assertRaises(TypeError, zipimport.zipimporter, TESTMOD, kwd=None) def testFilenameTooLong(self): self.assertZipFailure('A' * 33000) def testEmptyFile(self): support.unlink(TESTMOD) support.create_empty_file(TESTMOD) self.assertZipFailure(TESTMOD) def testFileUnreadable(self): support.unlink(TESTMOD) fd = os.open(TESTMOD, os.O_CREAT, 000) try: os.close(fd) self.assertZipFailure(TESTMOD) finally: # If we leave "the read-only bit" set on Windows, nothing can # delete TESTMOD, and later tests suffer bogus failures. os.chmod(TESTMOD, 0o666) support.unlink(TESTMOD) def testNotZipFile(self): support.unlink(TESTMOD) fp = open(TESTMOD, 'w+') fp.write('a' * 22) fp.close() self.assertZipFailure(TESTMOD) # XXX: disabled until this works on Big-endian machines def _testBogusZipFile(self): support.unlink(TESTMOD) fp = open(TESTMOD, 'w+') fp.write(struct.pack('=I', 0x06054B50)) fp.write('a' * 18) fp.close() z = zipimport.zipimporter(TESTMOD) try: self.assertRaises(TypeError, z.find_module, None) self.assertRaises(TypeError, z.load_module, None) self.assertRaises(TypeError, z.is_package, None) self.assertRaises(TypeError, z.get_code, None) self.assertRaises(TypeError, z.get_data, None) self.assertRaises(TypeError, z.get_source, None) error = zipimport.ZipImportError self.assertEqual(z.find_module('abc'), None) self.assertRaises(error, z.load_module, 'abc') self.assertRaises(error, z.get_code, 'abc') self.assertRaises(OSError, z.get_data, 'abc') self.assertRaises(error, z.get_source, 'abc') self.assertRaises(error, z.is_package, 'abc') finally: zipimport._zip_directory_cache.clear() def test_main(): try: support.run_unittest( UncompressedZipImportTestCase, CompressedZipImportTestCase, BadFileZipImportTestCase, ZipFileModifiedAfterImportTestCase, ) finally: support.unlink(TESTMOD) if __name__ == "__main__": test_main()