14adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport contextlib 24adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport imp 34adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport importlib 44adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport sys 54adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport unittest 64adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 74adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 84adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao@contextlib.contextmanager 94adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaodef uncache(*names): 104adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """Uncache a module from sys.modules. 114adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 124adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao A basic sanity check is performed to prevent uncaching modules that either 134adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao cannot/shouldn't be uncached. 144adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 154adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """ 164adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao for name in names: 174adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if name in ('sys', 'marshal', 'imp'): 184adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao raise ValueError( 194adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao "cannot uncache {0} as it will break _importlib".format(name)) 204adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 214adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao del sys.modules[name] 224adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao except KeyError: 234adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao pass 244adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 254adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao yield 264adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao finally: 274adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao for name in names: 284adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 294adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao del sys.modules[name] 304adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao except KeyError: 314adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao pass 324adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 334adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 344adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao@contextlib.contextmanager 354adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaodef import_state(**kwargs): 364adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """Context manager to manage the various importers and stored state in the 374adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao sys module. 384adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 394adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao The 'modules' attribute is not supported as the interpreter state stores a 404adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao pointer to the dict that the interpreter uses internally; 414adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao reassigning to sys.modules does not have the desired effect. 424adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 434adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """ 444adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao originals = {} 454adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 464adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao for attr, default in (('meta_path', []), ('path', []), 474adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao ('path_hooks', []), 484adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao ('path_importer_cache', {})): 494adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao originals[attr] = getattr(sys, attr) 504adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if attr in kwargs: 514adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao new_value = kwargs[attr] 524adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao del kwargs[attr] 534adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao else: 544adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao new_value = default 554adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao setattr(sys, attr, new_value) 564adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if len(kwargs): 574adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao raise ValueError( 584adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 'unrecognized arguments: {0}'.format(kwargs.keys())) 594adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao yield 604adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao finally: 614adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao for attr, value in originals.items(): 624adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao setattr(sys, attr, value) 634adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 644adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 654adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoclass mock_modules(object): 664adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 674adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """A mock importer/loader.""" 684adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 694adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __init__(self, *names): 704adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.modules = {} 714adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao for name in names: 724adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if not name.endswith('.__init__'): 734adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao import_name = name 744adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao else: 754adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao import_name = name[:-len('.__init__')] 764adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if '.' not in name: 774adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao package = None 784adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao elif import_name == name: 794adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao package = name.rsplit('.', 1)[0] 804adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao else: 814adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao package = import_name 824adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module = imp.new_module(import_name) 834adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module.__loader__ = self 844adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module.__file__ = '<mock __file__>' 854adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module.__package__ = package 864adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module.attr = name 874adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if import_name != name: 884adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module.__path__ = ['<mock __path__>'] 894adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.modules[import_name] = module 904adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 914adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __getitem__(self, name): 924adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return self.modules[name] 934adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 944adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def find_module(self, fullname, path=None): 954adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if fullname not in self.modules: 964adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return None 974adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao else: 984adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return self 994adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1004adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def load_module(self, fullname): 1014adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if fullname not in self.modules: 1024adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao raise ImportError 1034adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao else: 1044adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao sys.modules[fullname] = self.modules[fullname] 1054adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return self.modules[fullname] 1064adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1074adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __enter__(self): 1084adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self._uncache = uncache(*self.modules.keys()) 1094adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self._uncache.__enter__() 1104adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return self 1114adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1124adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __exit__(self, *exc_info): 1134adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self._uncache.__exit__(None, None, None) 1144adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1154adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1164adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1174adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoclass ImportModuleTests(unittest.TestCase): 1184adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1194adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """Test importlib.import_module.""" 1204adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1214adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def test_module_import(self): 1224adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # Test importing a top-level module. 1234adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with mock_modules('top_level') as mock: 1244adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with import_state(meta_path=[mock]): 1254adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module = importlib.import_module('top_level') 1264adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.assertEqual(module.__name__, 'top_level') 1274adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1284adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def test_absolute_package_import(self): 1294adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # Test importing a module from a package with an absolute name. 1304adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao pkg_name = 'pkg' 1314adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao pkg_long_name = '{0}.__init__'.format(pkg_name) 1324adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao name = '{0}.mod'.format(pkg_name) 1334adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with mock_modules(pkg_long_name, name) as mock: 1344adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with import_state(meta_path=[mock]): 1354adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module = importlib.import_module(name) 1364adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.assertEqual(module.__name__, name) 1374adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1384adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def test_shallow_relative_package_import(self): 1394adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao modules = ['a.__init__', 'a.b.__init__', 'a.b.c.__init__', 'a.b.c.d'] 1404adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with mock_modules(*modules) as mock: 1414adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with import_state(meta_path=[mock]): 1424adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module = importlib.import_module('.d', 'a.b.c') 1434adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.assertEqual(module.__name__, 'a.b.c.d') 1444adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1454adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def test_deep_relative_package_import(self): 1464adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # Test importing a module from a package through a relatve import. 1474adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao modules = ['a.__init__', 'a.b.__init__', 'a.c'] 1484adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with mock_modules(*modules) as mock: 1494adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with import_state(meta_path=[mock]): 1504adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module = importlib.import_module('..c', 'a.b') 1514adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.assertEqual(module.__name__, 'a.c') 1524adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1534adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def test_absolute_import_with_package(self): 1544adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # Test importing a module from a package with an absolute name with 1554adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # the 'package' argument given. 1564adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao pkg_name = 'pkg' 1574adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao pkg_long_name = '{0}.__init__'.format(pkg_name) 1584adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao name = '{0}.mod'.format(pkg_name) 1594adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with mock_modules(pkg_long_name, name) as mock: 1604adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with import_state(meta_path=[mock]): 1614adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao module = importlib.import_module(name, pkg_name) 1624adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.assertEqual(module.__name__, name) 1634adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1644adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def test_relative_import_wo_package(self): 1654adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # Relative imports cannot happen without the 'package' argument being 1664adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # set. 1674adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.assertRaises(TypeError, importlib.import_module, '.support') 1684adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1694adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1704adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaodef test_main(): 1714adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao from test.test_support import run_unittest 1724adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao run_unittest(ImportModuleTests) 1734adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1744adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1754adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoif __name__ == '__main__': 1764adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao test_main() 177