import sys import imp import os import unittest from test import test_support test_src = """\ def get_name(): return __name__ def get_file(): return __file__ """ absimp = "import sub\n" relimp = "from . import sub\n" deeprelimp = "from .... import sub\n" futimp = "from __future__ import absolute_import\n" reload_src = test_src+"""\ reloaded = True """ test_co = compile(test_src, "", "exec") reload_co = compile(reload_src, "", "exec") test2_oldabs_co = compile(absimp + test_src, "", "exec") test2_newabs_co = compile(futimp + absimp + test_src, "", "exec") test2_newrel_co = compile(relimp + test_src, "", "exec") test2_deeprel_co = compile(deeprelimp + test_src, "", "exec") test2_futrel_co = compile(futimp + relimp + test_src, "", "exec") test_path = "!!!_test_!!!" class ImportTracker: """Importer that only tracks attempted imports.""" def __init__(self): self.imports = [] def find_module(self, fullname, path=None): self.imports.append(fullname) return None class TestImporter: modules = { "hooktestmodule": (False, test_co), "hooktestpackage": (True, test_co), "hooktestpackage.sub": (True, test_co), "hooktestpackage.sub.subber": (True, test_co), "hooktestpackage.oldabs": (False, test2_oldabs_co), "hooktestpackage.newabs": (False, test2_newabs_co), "hooktestpackage.newrel": (False, test2_newrel_co), "hooktestpackage.sub.subber.subest": (True, test2_deeprel_co), "hooktestpackage.futrel": (False, test2_futrel_co), "sub": (False, test_co), "reloadmodule": (False, test_co), } def __init__(self, path=test_path): if path != test_path: # if out class is on sys.path_hooks, we must raise # ImportError for any path item that we can't handle. raise ImportError self.path = path def _get__path__(self): raise NotImplementedError def find_module(self, fullname, path=None): if fullname in self.modules: return self else: return None def load_module(self, fullname): ispkg, code = self.modules[fullname] mod = sys.modules.setdefault(fullname,imp.new_module(fullname)) mod.__file__ = "<%s>" % self.__class__.__name__ mod.__loader__ = self if ispkg: mod.__path__ = self._get__path__() exec code in mod.__dict__ return mod class MetaImporter(TestImporter): def _get__path__(self): return [] class PathImporter(TestImporter): def _get__path__(self): return [self.path] class ImportBlocker: """Place an ImportBlocker instance on sys.meta_path and you can be sure the modules you specified can't be imported, even if it's a builtin.""" def __init__(self, *namestoblock): self.namestoblock = dict.fromkeys(namestoblock) def find_module(self, fullname, path=None): if fullname in self.namestoblock: return self return None def load_module(self, fullname): raise ImportError, "I dare you" class ImpWrapper: def __init__(self, path=None): if path is not None and not os.path.isdir(path): raise ImportError self.path = path def find_module(self, fullname, path=None): subname = fullname.split(".")[-1] if subname != fullname and self.path is None: return None if self.path is None: path = None else: path = [self.path] try: file, filename, stuff = imp.find_module(subname, path) except ImportError: return None return ImpLoader(file, filename, stuff) class ImpLoader: def __init__(self, file, filename, stuff): self.file = file self.filename = filename self.stuff = stuff def load_module(self, fullname): mod = imp.load_module(fullname, self.file, self.filename, self.stuff) if self.file: self.file.close() mod.__loader__ = self # for introspection return mod class ImportHooksBaseTestCase(unittest.TestCase): def setUp(self): self.path = sys.path[:] self.meta_path = sys.meta_path[:] self.path_hooks = sys.path_hooks[:] sys.path_importer_cache.clear() self.tracker = ImportTracker() sys.meta_path.insert(0, self.tracker) def tearDown(self): sys.path[:] = self.path sys.meta_path[:] = self.meta_path sys.path_hooks[:] = self.path_hooks sys.path_importer_cache.clear() for fullname in self.tracker.imports: if fullname in sys.modules: del sys.modules[fullname] class ImportHooksTestCase(ImportHooksBaseTestCase): def doTestImports(self, importer=None): import hooktestmodule import hooktestpackage import hooktestpackage.sub import hooktestpackage.sub.subber self.assertEqual(hooktestmodule.get_name(), "hooktestmodule") self.assertEqual(hooktestpackage.get_name(), "hooktestpackage") self.assertEqual(hooktestpackage.sub.get_name(), "hooktestpackage.sub") self.assertEqual(hooktestpackage.sub.subber.get_name(), "hooktestpackage.sub.subber") if importer: self.assertEqual(hooktestmodule.__loader__, importer) self.assertEqual(hooktestpackage.__loader__, importer) self.assertEqual(hooktestpackage.sub.__loader__, importer) self.assertEqual(hooktestpackage.sub.subber.__loader__, importer) TestImporter.modules['reloadmodule'] = (False, test_co) import reloadmodule self.failIf(hasattr(reloadmodule,'reloaded')) TestImporter.modules['reloadmodule'] = (False, reload_co) reload(reloadmodule) self.failUnless(hasattr(reloadmodule,'reloaded')) import hooktestpackage.oldabs self.assertEqual(hooktestpackage.oldabs.get_name(), "hooktestpackage.oldabs") self.assertEqual(hooktestpackage.oldabs.sub, hooktestpackage.sub) import hooktestpackage.newrel self.assertEqual(hooktestpackage.newrel.get_name(), "hooktestpackage.newrel") self.assertEqual(hooktestpackage.newrel.sub, hooktestpackage.sub) import hooktestpackage.sub.subber.subest as subest self.assertEqual(subest.get_name(), "hooktestpackage.sub.subber.subest") self.assertEqual(subest.sub, hooktestpackage.sub) import hooktestpackage.futrel self.assertEqual(hooktestpackage.futrel.get_name(), "hooktestpackage.futrel") self.assertEqual(hooktestpackage.futrel.sub, hooktestpackage.sub) import sub self.assertEqual(sub.get_name(), "sub") import hooktestpackage.newabs self.assertEqual(hooktestpackage.newabs.get_name(), "hooktestpackage.newabs") self.assertEqual(hooktestpackage.newabs.sub, sub) def testMetaPath(self): i = MetaImporter() sys.meta_path.append(i) self.doTestImports(i) def testPathHook(self): sys.path_hooks.append(PathImporter) sys.path.append(test_path) self.doTestImports() def testBlocker(self): mname = "exceptions" # an arbitrary harmless builtin module if mname in sys.modules: del sys.modules[mname] sys.meta_path.append(ImportBlocker(mname)) try: __import__(mname) except ImportError: pass else: self.fail("'%s' was not supposed to be importable" % mname) def testImpWrapper(self): i = ImpWrapper() sys.meta_path.append(i) sys.path_hooks.append(ImpWrapper) mnames = ("colorsys", "urlparse", "distutils.core", "compiler.misc") for mname in mnames: parent = mname.split(".")[0] for n in sys.modules.keys(): if n.startswith(parent): del sys.modules[n] for mname in mnames: m = __import__(mname, globals(), locals(), ["__dummy__"]) m.__loader__ # to make sure we actually handled the import # Delete urllib from modules because urlparse was imported above. # Without this hack, test_socket_ssl fails if run in this order: # regrtest.py test_codecmaps_tw test_importhooks test_socket_ssl try: del sys.modules['urllib'] except KeyError: pass def test_main(): test_support.run_unittest(ImportHooksTestCase) if __name__ == "__main__": test_main()