importer.py 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. """This module implements a post import hook mechanism styled after what is
  2. described in PEP-369. Note that it doesn't cope with modules being reloaded.
  3. """
  4. import sys
  5. import threading
  6. PY2 = sys.version_info[0] == 2
  7. if PY2:
  8. string_types = basestring,
  9. else:
  10. import importlib
  11. string_types = str,
  12. from .decorators import synchronized
  13. # The dictionary registering any post import hooks to be triggered once
  14. # the target module has been imported. Once a module has been imported
  15. # and the hooks fired, the list of hooks recorded against the target
  16. # module will be truncated but the list left in the dictionary. This
  17. # acts as a flag to indicate that the module had already been imported.
  18. _post_import_hooks = {}
  19. _post_import_hooks_init = False
  20. _post_import_hooks_lock = threading.RLock()
  21. # Register a new post import hook for the target module name. This
  22. # differs from the PEP-369 implementation in that it also allows the
  23. # hook function to be specified as a string consisting of the name of
  24. # the callback in the form 'module:function'. This will result in a
  25. # proxy callback being registered which will defer loading of the
  26. # specified module containing the callback function until required.
  27. def _create_import_hook_from_string(name):
  28. def import_hook(module):
  29. module_name, function = name.split(':')
  30. attrs = function.split('.')
  31. __import__(module_name)
  32. callback = sys.modules[module_name]
  33. for attr in attrs:
  34. callback = getattr(callback, attr)
  35. return callback(module)
  36. return import_hook
  37. @synchronized(_post_import_hooks_lock)
  38. def register_post_import_hook(hook, name):
  39. # Create a deferred import hook if hook is a string name rather than
  40. # a callable function.
  41. if isinstance(hook, string_types):
  42. hook = _create_import_hook_from_string(hook)
  43. # Automatically install the import hook finder if it has not already
  44. # been installed.
  45. global _post_import_hooks_init
  46. if not _post_import_hooks_init:
  47. _post_import_hooks_init = True
  48. sys.meta_path.insert(0, ImportHookFinder())
  49. # Determine if any prior registration of a post import hook for
  50. # the target modules has occurred and act appropriately.
  51. hooks = _post_import_hooks.get(name, None)
  52. if hooks is None:
  53. # No prior registration of post import hooks for the target
  54. # module. We need to check whether the module has already been
  55. # imported. If it has we fire the hook immediately and add an
  56. # empty list to the registry to indicate that the module has
  57. # already been imported and hooks have fired. Otherwise add
  58. # the post import hook to the registry.
  59. module = sys.modules.get(name, None)
  60. if module is not None:
  61. _post_import_hooks[name] = []
  62. hook(module)
  63. else:
  64. _post_import_hooks[name] = [hook]
  65. elif hooks == []:
  66. # A prior registration of port import hooks for the target
  67. # module was done and the hooks already fired. Fire the hook
  68. # immediately.
  69. module = sys.modules[name]
  70. hook(module)
  71. else:
  72. # A prior registration of port import hooks for the target
  73. # module was done but the module has not yet been imported.
  74. _post_import_hooks[name].append(hook)
  75. # Register post import hooks defined as package entry points.
  76. def _create_import_hook_from_entrypoint(entrypoint):
  77. def import_hook(module):
  78. __import__(entrypoint.module_name)
  79. callback = sys.modules[entrypoint.module_name]
  80. for attr in entrypoint.attrs:
  81. callback = getattr(callback, attr)
  82. return callback(module)
  83. return import_hook
  84. def discover_post_import_hooks(group):
  85. try:
  86. import pkg_resources
  87. except ImportError:
  88. return
  89. for entrypoint in pkg_resources.iter_entry_points(group=group):
  90. callback = _create_import_hook_from_entrypoint(entrypoint)
  91. register_post_import_hook(callback, entrypoint.name)
  92. # Indicate that a module has been loaded. Any post import hooks which
  93. # were registered against the target module will be invoked. If an
  94. # exception is raised in any of the post import hooks, that will cause
  95. # the import of the target module to fail.
  96. @synchronized(_post_import_hooks_lock)
  97. def notify_module_loaded(module):
  98. name = getattr(module, '__name__', None)
  99. hooks = _post_import_hooks.get(name, None)
  100. if hooks:
  101. _post_import_hooks[name] = []
  102. for hook in hooks:
  103. hook(module)
  104. # A custom module import finder. This intercepts attempts to import
  105. # modules and watches out for attempts to import target modules of
  106. # interest. When a module of interest is imported, then any post import
  107. # hooks which are registered will be invoked.
  108. class _ImportHookLoader:
  109. def load_module(self, fullname):
  110. module = sys.modules[fullname]
  111. notify_module_loaded(module)
  112. return module
  113. class _ImportHookChainedLoader:
  114. def __init__(self, loader):
  115. self.loader = loader
  116. def load_module(self, fullname):
  117. module = self.loader.load_module(fullname)
  118. notify_module_loaded(module)
  119. return module
  120. class ImportHookFinder:
  121. def __init__(self):
  122. self.in_progress = {}
  123. @synchronized(_post_import_hooks_lock)
  124. def find_module(self, fullname, path=None):
  125. # If the module being imported is not one we have registered
  126. # post import hooks for, we can return immediately. We will
  127. # take no further part in the importing of this module.
  128. if not fullname in _post_import_hooks:
  129. return None
  130. # When we are interested in a specific module, we will call back
  131. # into the import system a second time to defer to the import
  132. # finder that is supposed to handle the importing of the module.
  133. # We set an in progress flag for the target module so that on
  134. # the second time through we don't trigger another call back
  135. # into the import system and cause a infinite loop.
  136. if fullname in self.in_progress:
  137. return None
  138. self.in_progress[fullname] = True
  139. # Now call back into the import system again.
  140. try:
  141. if PY2:
  142. # For Python 2 we don't have much choice but to
  143. # call back in to __import__(). This will
  144. # actually cause the module to be imported. If no
  145. # module could be found then ImportError will be
  146. # raised. Otherwise we return a loader which
  147. # returns the already loaded module and invokes
  148. # the post import hooks.
  149. __import__(fullname)
  150. return _ImportHookLoader()
  151. else:
  152. # For Python 3 we need to use find_spec().loader
  153. # from the importlib.util module. It doesn't actually
  154. # import the target module and only finds the
  155. # loader. If a loader is found, we need to return
  156. # our own loader which will then in turn call the
  157. # real loader to import the module and invoke the
  158. # post import hooks.
  159. try:
  160. import importlib.util
  161. loader = importlib.util.find_spec(fullname).loader
  162. except (ImportError, AttributeError):
  163. loader = importlib.find_loader(fullname, path)
  164. if loader:
  165. return _ImportHookChainedLoader(loader)
  166. finally:
  167. del self.in_progress[fullname]
  168. # Decorator for marking that a function should be called as a post
  169. # import hook when the target module is imported.
  170. def when_imported(name):
  171. def register(hook):
  172. register_post_import_hook(hook, name)
  173. return hook
  174. return register