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