10c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport contextlib 20c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport imp 30c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport importlib 40c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport sys 50c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport unittest 60c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 70c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 80c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi@contextlib.contextmanager 90c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef uncache(*names): 100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi """Uncache a module from sys.modules. 110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi A basic sanity check is performed to prevent uncaching modules that either 130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi cannot/shouldn't be uncached. 140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi """ 160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for name in names: 170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if name in ('sys', 'marshal', 'imp'): 180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi raise ValueError( 190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi "cannot uncache {0} as it will break _importlib".format(name)) 200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi try: 210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi del sys.modules[name] 220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi except KeyError: 230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pass 240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi try: 250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi yield 260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi finally: 270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for name in names: 280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi try: 290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi del sys.modules[name] 300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi except KeyError: 310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pass 320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi@contextlib.contextmanager 350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef import_state(**kwargs): 360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi """Context manager to manage the various importers and stored state in the 370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi sys module. 380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi The 'modules' attribute is not supported as the interpreter state stores a 400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pointer to the dict that the interpreter uses internally; 410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi reassigning to sys.modules does not have the desired effect. 420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi """ 440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi originals = {} 450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi try: 460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for attr, default in (('meta_path', []), ('path', []), 470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi ('path_hooks', []), 480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi ('path_importer_cache', {})): 490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi originals[attr] = getattr(sys, attr) 500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if attr in kwargs: 510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi new_value = kwargs[attr] 520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi del kwargs[attr] 530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi else: 540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi new_value = default 550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi setattr(sys, attr, new_value) 560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if len(kwargs): 570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi raise ValueError( 580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'unrecognized arguments: {0}'.format(kwargs.keys())) 590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi yield 600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi finally: 610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for attr, value in originals.items(): 620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi setattr(sys, attr, value) 630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass mock_modules(object): 660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi """A mock importer/loader.""" 680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, *names): 700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.modules = {} 710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for name in names: 720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if not name.endswith('.__init__'): 730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi import_name = name 740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi else: 750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi import_name = name[:-len('.__init__')] 760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if '.' not in name: 770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi package = None 780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi elif import_name == name: 790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi package = name.rsplit('.', 1)[0] 800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi else: 810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi package = import_name 820c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = imp.new_module(import_name) 830c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module.__loader__ = self 840c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module.__file__ = '<mock __file__>' 850c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module.__package__ = package 860c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module.attr = name 870c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if import_name != name: 880c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module.__path__ = ['<mock __path__>'] 890c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.modules[import_name] = module 900c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 910c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __getitem__(self, name): 920c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self.modules[name] 930c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 940c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def find_module(self, fullname, path=None): 950c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if fullname not in self.modules: 960c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return None 970c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi else: 980c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self 990c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1000c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def load_module(self, fullname): 1010c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if fullname not in self.modules: 1020c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi raise ImportError 1030c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi else: 1040c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi sys.modules[fullname] = self.modules[fullname] 1050c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self.modules[fullname] 1060c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1070c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __enter__(self): 1080c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self._uncache = uncache(*self.modules.keys()) 1090c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self._uncache.__enter__() 1100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self 1110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __exit__(self, *exc_info): 1130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self._uncache.__exit__(None, None, None) 1140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass ImportModuleTests(unittest.TestCase): 1180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi """Test importlib.import_module.""" 1200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_module_import(self): 1220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Test importing a top-level module. 1230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with mock_modules('top_level') as mock: 1240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with import_state(meta_path=[mock]): 1250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = importlib.import_module('top_level') 1260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(module.__name__, 'top_level') 1270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_absolute_package_import(self): 1290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Test importing a module from a package with an absolute name. 1300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pkg_name = 'pkg' 1310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pkg_long_name = '{0}.__init__'.format(pkg_name) 1320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi name = '{0}.mod'.format(pkg_name) 1330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with mock_modules(pkg_long_name, name) as mock: 1340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with import_state(meta_path=[mock]): 1350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = importlib.import_module(name) 1360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(module.__name__, name) 1370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_shallow_relative_package_import(self): 1390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi modules = ['a.__init__', 'a.b.__init__', 'a.b.c.__init__', 'a.b.c.d'] 1400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with mock_modules(*modules) as mock: 1410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with import_state(meta_path=[mock]): 1420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = importlib.import_module('.d', 'a.b.c') 1430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(module.__name__, 'a.b.c.d') 1440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_deep_relative_package_import(self): 1460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Test importing a module from a package through a relatve import. 1470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi modules = ['a.__init__', 'a.b.__init__', 'a.c'] 1480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with mock_modules(*modules) as mock: 1490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with import_state(meta_path=[mock]): 1500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = importlib.import_module('..c', 'a.b') 1510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(module.__name__, 'a.c') 1520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_absolute_import_with_package(self): 1540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Test importing a module from a package with an absolute name with 1550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # the 'package' argument given. 1560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pkg_name = 'pkg' 1570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pkg_long_name = '{0}.__init__'.format(pkg_name) 1580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi name = '{0}.mod'.format(pkg_name) 1590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with mock_modules(pkg_long_name, name) as mock: 1600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with import_state(meta_path=[mock]): 1610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = importlib.import_module(name, pkg_name) 1620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(module.__name__, name) 1630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_relative_import_wo_package(self): 1650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Relative imports cannot happen without the 'package' argument being 1660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # set. 1670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, importlib.import_module, '.support') 1680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef test_main(): 1710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi from test.test_support import run_unittest 1720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi run_unittest(ImportModuleTests) 1730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiif __name__ == '__main__': 1760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi test_main() 177