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