import sys import os import marshal import imp import struct import time import unittest from test import support from test.test_importhooks import ImportHooksBaseTestCase, test_src, test_co # some tests can be ran even without zlib try: import zlib except ImportError: zlib = None 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 raise_src = 'def do_raise(): raise TypeError\n' def make_pyc(co, mtime): 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 = imp.get_magic() + struct.pack("", "exec"), NOW) 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) 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) @unittest.skipUnless(zlib, "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) def testFilenameTooLong(self): self.assertZipFailure('A' * 33000) def testEmptyFile(self): support.unlink(TESTMOD) open(TESTMOD, 'w+').close() 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(IOError, 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: support.unlink(TESTMOD) if __name__ == "__main__": test_main()