importer.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293
  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. find_spec = None
  10. else:
  11. string_types = str,
  12. from importlib.util import find_spec
  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. def register_post_import_hook(hook, name):
  38. # Create a deferred import hook if hook is a string name rather than
  39. # a callable function.
  40. if isinstance(hook, string_types):
  41. hook = _create_import_hook_from_string(hook)
  42. with _post_import_hooks_lock:
  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. # Check if the module is already imported. If not, register the hook
  50. # to be called after import.
  51. module = sys.modules.get(name, None)
  52. if module is None:
  53. _post_import_hooks.setdefault(name, []).append(hook)
  54. # If the module is already imported, we fire the hook right away. Note that
  55. # the hook is called outside of the lock to avoid deadlocks if code run as a
  56. # consequence of calling the module import hook in turn triggers a separate
  57. # thread which tries to register an import hook.
  58. if module is not None:
  59. hook(module)
  60. # Register post import hooks defined as package entry points.
  61. def _create_import_hook_from_entrypoint(entrypoint):
  62. def import_hook(module):
  63. __import__(entrypoint.module_name)
  64. callback = sys.modules[entrypoint.module_name]
  65. for attr in entrypoint.attrs:
  66. callback = getattr(callback, attr)
  67. return callback(module)
  68. return import_hook
  69. def discover_post_import_hooks(group):
  70. try:
  71. import pkg_resources
  72. except ImportError:
  73. return
  74. for entrypoint in pkg_resources.iter_entry_points(group=group):
  75. callback = _create_import_hook_from_entrypoint(entrypoint)
  76. register_post_import_hook(callback, entrypoint.name)
  77. # Indicate that a module has been loaded. Any post import hooks which
  78. # were registered against the target module will be invoked. If an
  79. # exception is raised in any of the post import hooks, that will cause
  80. # the import of the target module to fail.
  81. def notify_module_loaded(module):
  82. name = getattr(module, '__name__', None)
  83. with _post_import_hooks_lock:
  84. hooks = _post_import_hooks.pop(name, ())
  85. # Note that the hook is called outside of the lock to avoid deadlocks if
  86. # code run as a consequence of calling the module import hook in turn
  87. # triggers a separate thread which tries to register an import hook.
  88. for hook in hooks:
  89. hook(module)
  90. # A custom module import finder. This intercepts attempts to import
  91. # modules and watches out for attempts to import target modules of
  92. # interest. When a module of interest is imported, then any post import
  93. # hooks which are registered will be invoked.
  94. class _ImportHookLoader:
  95. def load_module(self, fullname):
  96. module = sys.modules[fullname]
  97. notify_module_loaded(module)
  98. return module
  99. class _ImportHookChainedLoader:
  100. def __init__(self, loader):
  101. self.loader = loader
  102. if hasattr(loader, "load_module"):
  103. self.load_module = self._load_module
  104. if hasattr(loader, "create_module"):
  105. self.create_module = self._create_module
  106. if hasattr(loader, "exec_module"):
  107. self.exec_module = self._exec_module
  108. def _set_loader(self, module):
  109. # Set module's loader to self.loader unless it's already set to
  110. # something else. Import machinery will set it to spec.loader if it is
  111. # None, so handle None as well. The module may not support attribute
  112. # assignment, in which case we simply skip it. Note that we also deal
  113. # with __loader__ not existing at all. This is to future proof things
  114. # due to proposal to remove the attribue as described in the GitHub
  115. # issue at https://github.com/python/cpython/issues/77458. Also prior
  116. # to Python 3.3, the __loader__ attribute was only set if a custom
  117. # module loader was used. It isn't clear whether the attribute still
  118. # existed in that case or was set to None.
  119. class UNDEFINED: pass
  120. if getattr(module, "__loader__", UNDEFINED) in (None, self):
  121. try:
  122. module.__loader__ = self.loader
  123. except AttributeError:
  124. pass
  125. if (getattr(module, "__spec__", None) is not None
  126. and getattr(module.__spec__, "loader", None) is self):
  127. module.__spec__.loader = self.loader
  128. def _load_module(self, fullname):
  129. module = self.loader.load_module(fullname)
  130. self._set_loader(module)
  131. notify_module_loaded(module)
  132. return module
  133. # Python 3.4 introduced create_module() and exec_module() instead of
  134. # load_module() alone. Splitting the two steps.
  135. def _create_module(self, spec):
  136. return self.loader.create_module(spec)
  137. def _exec_module(self, module):
  138. self._set_loader(module)
  139. self.loader.exec_module(module)
  140. notify_module_loaded(module)
  141. class ImportHookFinder:
  142. def __init__(self):
  143. self.in_progress = {}
  144. def find_module(self, fullname, path=None):
  145. # If the module being imported is not one we have registered
  146. # post import hooks for, we can return immediately. We will
  147. # take no further part in the importing of this module.
  148. with _post_import_hooks_lock:
  149. if fullname not in _post_import_hooks:
  150. return None
  151. # When we are interested in a specific module, we will call back
  152. # into the import system a second time to defer to the import
  153. # finder that is supposed to handle the importing of the module.
  154. # We set an in progress flag for the target module so that on
  155. # the second time through we don't trigger another call back
  156. # into the import system and cause a infinite loop.
  157. if fullname in self.in_progress:
  158. return None
  159. self.in_progress[fullname] = True
  160. # Now call back into the import system again.
  161. try:
  162. if not find_spec:
  163. # For Python 2 we don't have much choice but to
  164. # call back in to __import__(). This will
  165. # actually cause the module to be imported. If no
  166. # module could be found then ImportError will be
  167. # raised. Otherwise we return a loader which
  168. # returns the already loaded module and invokes
  169. # the post import hooks.
  170. __import__(fullname)
  171. return _ImportHookLoader()
  172. else:
  173. # For Python 3 we need to use find_spec().loader
  174. # from the importlib.util module. It doesn't actually
  175. # import the target module and only finds the
  176. # loader. If a loader is found, we need to return
  177. # our own loader which will then in turn call the
  178. # real loader to import the module and invoke the
  179. # post import hooks.
  180. loader = getattr(find_spec(fullname), "loader", None)
  181. if loader and not isinstance(loader, _ImportHookChainedLoader):
  182. return _ImportHookChainedLoader(loader)
  183. finally:
  184. del self.in_progress[fullname]
  185. def find_spec(self, fullname, path=None, target=None):
  186. # Since Python 3.4, you are meant to implement find_spec() method
  187. # instead of find_module() and since Python 3.10 you get deprecation
  188. # warnings if you don't define find_spec().
  189. # If the module being imported is not one we have registered
  190. # post import hooks for, we can return immediately. We will
  191. # take no further part in the importing of this module.
  192. with _post_import_hooks_lock:
  193. if fullname not in _post_import_hooks:
  194. return None
  195. # When we are interested in a specific module, we will call back
  196. # into the import system a second time to defer to the import
  197. # finder that is supposed to handle the importing of the module.
  198. # We set an in progress flag for the target module so that on
  199. # the second time through we don't trigger another call back
  200. # into the import system and cause a infinite loop.
  201. if fullname in self.in_progress:
  202. return None
  203. self.in_progress[fullname] = True
  204. # Now call back into the import system again.
  205. try:
  206. # This should only be Python 3 so find_spec() should always
  207. # exist so don't need to check.
  208. spec = find_spec(fullname)
  209. loader = getattr(spec, "loader", None)
  210. if loader and not isinstance(loader, _ImportHookChainedLoader):
  211. spec.loader = _ImportHookChainedLoader(loader)
  212. return spec
  213. finally:
  214. del self.in_progress[fullname]
  215. # Decorator for marking that a function should be called as a post
  216. # import hook when the target module is imported.
  217. def when_imported(name):
  218. def register(hook):
  219. register_post_import_hook(hook, name)
  220. return hook
  221. return register