Lib/test/test_importhooks.py

   1 import sys
   2 import imp
   3 import os
   4 import unittest
   5 from test import test_support
   6
   7
   8 test_src = """\
   9 def get_name():
  10     return __name__
  11 def get_file():
  12     return __file__
  13 """
  14
  15 absimp = "import sub\n"
  16 relimp = "from . import sub\n"
  17 deeprelimp = "from .... import sub\n"
  18 futimp = "from __future__ import absolute_import\n"
  19
  20 reload_src = test_src+"""\
  21 reloaded = True
  22 """
  23
  24 test_co = compile(test_src, "<???>", "exec")
  25 reload_co = compile(reload_src, "<???>", "exec")
  26
  27 test2_oldabs_co = compile(absimp + test_src, "<???>", "exec")
  28 test2_newabs_co = compile(futimp + absimp + test_src, "<???>", "exec")
  29 test2_newrel_co = compile(relimp + test_src, "<???>", "exec")
  30 test2_deeprel_co = compile(deeprelimp + test_src, "<???>", "exec")
  31 test2_futrel_co = compile(futimp + relimp + test_src, "<???>", "exec")
  32
  33 test_path = "!!!_test_!!!"
  34
  35
  36 class ImportTracker:
  37     """Importer that only tracks attempted imports."""
  38     def __init__(self):
  39         self.imports = []
  40     def find_module(self, fullname, path=None):
  41         self.imports.append(fullname)
  42         return None
  43
  44
  45 class TestImporter:
  46
  47     modules = {
  48         "hooktestmodule": (False, test_co),
  49         "hooktestpackage": (True, test_co),
  50         "hooktestpackage.sub": (True, test_co),
  51         "hooktestpackage.sub.subber": (True, test_co),
  52         "hooktestpackage.oldabs": (False, test2_oldabs_co),
  53         "hooktestpackage.newabs": (False, test2_newabs_co),
  54         "hooktestpackage.newrel": (False, test2_newrel_co),
  55         "hooktestpackage.sub.subber.subest": (True, test2_deeprel_co),
  56         "hooktestpackage.futrel": (False, test2_futrel_co),
  57         "sub": (False, test_co),
  58         "reloadmodule": (False, test_co),
  59     }
  60
  61     def __init__(self, path=test_path):
  62         if path != test_path:
  63             # if out class is on sys.path_hooks, we must raise
  64             # ImportError for any path item that we can't handle.
  65             raise ImportError
  66         self.path = path
  67
  68     def _get__path__(self):
  69         raise NotImplementedError
  70
  71     def find_module(self, fullname, path=None):
  72         if fullname in self.modules:
  73             return self
  74         else:
  75             return None
  76
  77     def load_module(self, fullname):
  78         ispkg, code = self.modules[fullname]
  79         mod = sys.modules.setdefault(fullname,imp.new_module(fullname))
  80         mod.__file__ = "<%s>" % self.__class__.__name__
  81         mod.__loader__ = self
  82         if ispkg:
  83             mod.__path__ = self._get__path__()
  84         exec code in mod.__dict__
  85         return mod
  86
  87
  88 class MetaImporter(TestImporter):
  89     def _get__path__(self):
  90         return []
  91
  92 class PathImporter(TestImporter):
  93     def _get__path__(self):
  94         return [self.path]
  95
  96
  97 class ImportBlocker:
  98     """Place an ImportBlocker instance on sys.meta_path and you
  99     can be sure the modules you specified can't be imported, even
 100     if it's a builtin."""
 101     def __init__(self, *namestoblock):
 102         self.namestoblock = dict.fromkeys(namestoblock)
 103     def find_module(self, fullname, path=None):
 104         if fullname in self.namestoblock:
 105             return self
 106         return None
 107     def load_module(self, fullname):
 108         raise ImportError, "I dare you"
 109
 110
 111 class ImpWrapper:
 112
 113     def __init__(self, path=None):
 114         if path is not None and not os.path.isdir(path):
 115             raise ImportError
 116         self.path = path
 117
 118     def find_module(self, fullname, path=None):
 119         subname = fullname.split(".")[-1]
 120         if subname != fullname and self.path is None:
 121             return None
 122         if self.path is None:
 123             path = None
 124         else:
 125             path = [self.path]
 126         try:
 127             file, filename, stuff = imp.find_module(subname, path)
 128         except ImportError:
 129             return None
 130         return ImpLoader(file, filename, stuff)
 131
 132
 133 class ImpLoader:
 134
 135     def __init__(self, file, filename, stuff):
 136         self.file = file
 137         self.filename = filename
 138         self.stuff = stuff
 139
 140     def load_module(self, fullname):
 141         mod = imp.load_module(fullname, self.file, self.filename, self.stuff)
 142         if self.file:
 143             self.file.close()
 144         mod.__loader__ = self  # for introspection
 145         return mod
 146
 147
 148 class ImportHooksBaseTestCase(unittest.TestCase):
 149
 150     def setUp(self):
 151         self.path = sys.path[:]
 152         self.meta_path = sys.meta_path[:]
 153         self.path_hooks = sys.path_hooks[:]
 154         sys.path_importer_cache.clear()
 155         self.tracker = ImportTracker()
 156         sys.meta_path.insert(0, self.tracker)
 157
 158     def tearDown(self):
 159         sys.path[:] = self.path
 160         sys.meta_path[:] = self.meta_path
 161         sys.path_hooks[:] = self.path_hooks
 162         sys.path_importer_cache.clear()
 163         for fullname in self.tracker.imports:
 164             if fullname in sys.modules:
 165                 del sys.modules[fullname]
 166
 167
 168 class ImportHooksTestCase(ImportHooksBaseTestCase):
 169
 170     def doTestImports(self, importer=None):
 171         import hooktestmodule
 172         import hooktestpackage
 173         import hooktestpackage.sub
 174         import hooktestpackage.sub.subber
 175         self.assertEqual(hooktestmodule.get_name(),
 176                          "hooktestmodule")
 177         self.assertEqual(hooktestpackage.get_name(),
 178                          "hooktestpackage")
 179         self.assertEqual(hooktestpackage.sub.get_name(),
 180                          "hooktestpackage.sub")
 181         self.assertEqual(hooktestpackage.sub.subber.get_name(),
 182                          "hooktestpackage.sub.subber")
 183         if importer:
 184             self.assertEqual(hooktestmodule.__loader__, importer)
 185             self.assertEqual(hooktestpackage.__loader__, importer)
 186             self.assertEqual(hooktestpackage.sub.__loader__, importer)
 187             self.assertEqual(hooktestpackage.sub.subber.__loader__, importer)
 188
 189         TestImporter.modules['reloadmodule'] = (False, test_co)
 190         import reloadmodule
 191         self.failIf(hasattr(reloadmodule,'reloaded'))
 192
 193         TestImporter.modules['reloadmodule'] = (False, reload_co)
 194         reload(reloadmodule)
 195         self.failUnless(hasattr(reloadmodule,'reloaded'))
 196
 197         import hooktestpackage.oldabs
 198         self.assertEqual(hooktestpackage.oldabs.get_name(),
 199                          "hooktestpackage.oldabs")
 200         self.assertEqual(hooktestpackage.oldabs.sub,
 201                          hooktestpackage.sub)
 202
 203         import hooktestpackage.newrel
 204         self.assertEqual(hooktestpackage.newrel.get_name(),
 205                          "hooktestpackage.newrel")
 206         self.assertEqual(hooktestpackage.newrel.sub,
 207                          hooktestpackage.sub)
 208
 209         import hooktestpackage.sub.subber.subest as subest
 210         self.assertEqual(subest.get_name(),
 211                          "hooktestpackage.sub.subber.subest")
 212         self.assertEqual(subest.sub,
 213                          hooktestpackage.sub)
 214
 215         import hooktestpackage.futrel
 216         self.assertEqual(hooktestpackage.futrel.get_name(),
 217                          "hooktestpackage.futrel")
 218         self.assertEqual(hooktestpackage.futrel.sub,
 219                          hooktestpackage.sub)
 220
 221         import sub
 222         self.assertEqual(sub.get_name(), "sub")
 223
 224         import hooktestpackage.newabs
 225         self.assertEqual(hooktestpackage.newabs.get_name(),
 226                          "hooktestpackage.newabs")
 227         self.assertEqual(hooktestpackage.newabs.sub, sub)
 228
 229     def testMetaPath(self):
 230         i = MetaImporter()
 231         sys.meta_path.append(i)
 232         self.doTestImports(i)
 233
 234     def testPathHook(self):
 235         sys.path_hooks.append(PathImporter)
 236         sys.path.append(test_path)
 237         self.doTestImports()
 238
 239     def testBlocker(self):
 240         mname = "exceptions"  # an arbitrary harmless builtin module
 241         if mname in sys.modules:
 242             del sys.modules[mname]
 243         sys.meta_path.append(ImportBlocker(mname))
 244         try:
 245             __import__(mname)
 246         except ImportError:
 247             pass
 248         else:
 249             self.fail("'%s' was not supposed to be importable" % mname)
 250
 251     def testImpWrapper(self):
 252         i = ImpWrapper()
 253         sys.meta_path.append(i)
 254         sys.path_hooks.append(ImpWrapper)
 255         mnames = ("colorsys", "urlparse", "distutils.core", "compiler.misc")
 256         for mname in mnames:
 257             parent = mname.split(".")[0]
 258             for n in sys.modules.keys():
 259                 if n.startswith(parent):
 260                     del sys.modules[n]
 261         for mname in mnames:
 262             m = __import__(mname, globals(), locals(), ["__dummy__"])
 263             m.__loader__  # to make sure we actually handled the import
 264         # Delete urllib from modules because urlparse was imported above.
 265         # Without this hack, test_socket_ssl fails if run in this order:
 266         # regrtest.py test_codecmaps_tw test_importhooks test_socket_ssl
 267         try:
 268             del sys.modules['urllib']
 269         except KeyError:
 270             pass
 271
 272 def test_main():
 273     test_support.run_unittest(ImportHooksTestCase)
 274
 275 if __name__ == "__main__":
 276     test_main()