import sys import os import marshal import importlib import importlib.util import struct import time import unittest import unittest.mock from test import support from test.support import import_helper from test.support import os_helper 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 try: import zlib except ImportError: zlib = None 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 testDefaultOptimizationLevel(self): # zipimport should use the default optimization level (#28131) src = """if 1: # indent hack def test(val): assert(val) return val\n""" files = {TESTMOD + '.py': (NOW, src)} self.makeZip(files) sys.path.insert(0, TEMP_ZIP) mod = importlib.import_module(TESTMOD) self.assertEqual(mod.test(1), 1) self.assertRaises(AssertionError, mod.test, False) 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(os_helper.TESTFN_UNENCODABLE is None, "need an unencodable filename") def testUnencodable(self): filename = os_helper.TESTFN_UNENCODABLE + ".zip" self.addCleanup(os_helper.unlink, filename) with ZipFile(filename, "w") as z: zinfo = ZipInfo(TESTMOD + ".py", time.localtime(NOW)) zinfo.compress_type = self.compression z.writestr(zinfo, test_src) spec = zipimport.zipimporter(filename).find_spec(TESTMOD) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) def testBytesPath(self): filename = os_helper.TESTFN + ".zip" self.addCleanup(os_helper.unlink, filename) with ZipFile(filename, "w") as z: zinfo = ZipInfo(TESTMOD + ".py", time.localtime(NOW)) zinfo.compress_type = self.compression z.writestr(zinfo, test_src) zipimport.zipimporter(filename) zipimport.zipimporter(os.fsencode(filename)) with self.assertRaises(TypeError): zipimport.zipimporter(bytearray(os.fsencode(filename))) with self.assertRaises(TypeError): zipimport.zipimporter(memoryview(os.fsencode(filename))) def testComment(self): files = {TESTMOD + ".py": (NOW, test_src)} self.doTest(".py", files, TESTMOD, comment=b"comment") def testBeginningCruftAndComment(self): files = {TESTMOD + ".py": (NOW, test_src)} self.doTest(".py", files, TESTMOD, stuff=b"cruft" * 64, comment=b"hi") def testLargestPossibleComment(self): files = {TESTMOD + ".py": (NOW, test_src)} self.doTest(".py", files, TESTMOD, comment=b"c" * ((1 << 16) - 1)) @support.requires_zlib() class CompressedZipImportTestCase(UncompressedZipImportTestCase): compression = ZIP_DEFLATED class BadFileZipImportTestCase(unittest.TestCase): def assertZipFailure(self, filename): 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) self.assertRaises(TypeError, zipimport.zipimporter, list(os.fsencode(TESTMOD))) def testFilenameTooLong(self): self.assertZipFailure('A' * 33000) def testEmptyFile(self): os_helper.unlink(TESTMOD) os_helper.create_empty_file(TESTMOD) self.assertZipFailure(TESTMOD) def testFileUnreadable(self): os_helper.unlink(TESTMOD) fd = os.open(TESTMOD, os.O_CREAT, 000) try: os.close(fd) with self.assertRaises(zipimport.ZipImportError) as cm: zipimport.zipimporter(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) os_helper.unlink(TESTMOD) def testNotZipFile(self): os_helper.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): os_helper.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.find_spec, None) self.assertRaises(TypeError, z.exec_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.assertIsNone(z.find_module('abc')) self.assertIsNone(z.find_spec('abc')) 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, ) finally: os_helper.unlink(TESTMOD) if __name__ == "__main__": test_main()