333 lines
12 KiB
Python
333 lines
12 KiB
Python
"""This module implements a post import hook mechanism styled after what is
|
|
described in PEP-369. Note that it doesn't cope with modules being reloaded.
|
|
|
|
"""
|
|
|
|
import importlib.metadata
|
|
import sys
|
|
import threading
|
|
from importlib.util import find_spec
|
|
from typing import Callable, Dict, List
|
|
|
|
from .__wrapt__ import BaseObjectProxy
|
|
|
|
# The dictionary registering any post import hooks to be triggered once
|
|
# the target module has been imported. Once a module has been imported
|
|
# and the hooks fired, the list of hooks recorded against the target
|
|
# module will be truncated but the list left in the dictionary. This
|
|
# acts as a flag to indicate that the module had already been imported.
|
|
|
|
_post_import_hooks: Dict[str, List[Callable]] = {}
|
|
_post_import_hooks_init = False
|
|
_post_import_hooks_lock = threading.RLock()
|
|
|
|
# Register a new post import hook for the target module name. This
|
|
# differs from the PEP-369 implementation in that it also allows the
|
|
# hook function to be specified as a string consisting of the name of
|
|
# the callback in the form 'module:function'. This will result in a
|
|
# proxy callback being registered which will defer loading of the
|
|
# specified module containing the callback function until required.
|
|
|
|
|
|
def _create_import_hook_from_string(name):
|
|
def import_hook(module):
|
|
module_name, function = name.split(":")
|
|
attrs = function.split(".")
|
|
__import__(module_name)
|
|
callback = sys.modules[module_name]
|
|
for attr in attrs:
|
|
callback = getattr(callback, attr)
|
|
return callback(module)
|
|
|
|
return import_hook
|
|
|
|
|
|
def register_post_import_hook(hook, name):
|
|
"""
|
|
Register a post import hook for the target module `name`. The `hook`
|
|
function will be called once the module is imported and will be passed the
|
|
module as argument. If the module is already imported, the `hook` will be
|
|
called immediately. If you also want to defer loading of the module containing
|
|
the `hook` function until required, you can specify the `hook` as a string in
|
|
the form 'module:function'. This will result in a proxy hook function being
|
|
registered which will defer loading of the specified module containing the
|
|
callback function until required.
|
|
"""
|
|
|
|
# Create a deferred import hook if hook is a string name rather than
|
|
# a callable function.
|
|
|
|
if isinstance(hook, str):
|
|
hook = _create_import_hook_from_string(hook)
|
|
|
|
with _post_import_hooks_lock:
|
|
# Automatically install the import hook finder if it has not already
|
|
# been installed.
|
|
|
|
global _post_import_hooks_init
|
|
|
|
if not _post_import_hooks_init:
|
|
_post_import_hooks_init = True
|
|
sys.meta_path.insert(0, ImportHookFinder())
|
|
|
|
# Check if the module is already imported. If not, register the hook
|
|
# to be called after import.
|
|
|
|
module = sys.modules.get(name, None)
|
|
|
|
if module is None:
|
|
_post_import_hooks.setdefault(name, []).append(hook)
|
|
|
|
# If the module is already imported, we fire the hook right away. Note that
|
|
# the hook is called outside of the lock to avoid deadlocks if code run as a
|
|
# consequence of calling the module import hook in turn triggers a separate
|
|
# thread which tries to register an import hook.
|
|
|
|
if module is not None:
|
|
hook(module)
|
|
|
|
|
|
# Register post import hooks defined as package entry points.
|
|
|
|
|
|
def _create_import_hook_from_entrypoint(entrypoint):
|
|
def import_hook(module):
|
|
entrypoint_value = entrypoint.value.split(":")
|
|
module_name = entrypoint_value[0]
|
|
__import__(module_name)
|
|
callback = sys.modules[module_name]
|
|
|
|
if len(entrypoint_value) > 1:
|
|
attrs = entrypoint_value[1].split(".")
|
|
for attr in attrs:
|
|
callback = getattr(callback, attr)
|
|
return callback(module)
|
|
|
|
return import_hook
|
|
|
|
|
|
def discover_post_import_hooks(group):
|
|
"""
|
|
Discover and register post import hooks defined as package entry points
|
|
in the specified `group`. The group should be a string that matches the
|
|
entry point group name used in the package metadata.
|
|
"""
|
|
|
|
try:
|
|
# Python 3.10+ style with select parameter
|
|
entrypoints = importlib.metadata.entry_points(group=group)
|
|
except TypeError:
|
|
# Python 3.8-3.9 style that returns a dict
|
|
entrypoints = importlib.metadata.entry_points().get(group, ())
|
|
|
|
for entrypoint in entrypoints:
|
|
callback = entrypoint.load() # Use the loaded callback directly
|
|
register_post_import_hook(callback, entrypoint.name)
|
|
|
|
|
|
# Indicate that a module has been loaded. Any post import hooks which
|
|
# were registered against the target module will be invoked. If an
|
|
# exception is raised in any of the post import hooks, that will cause
|
|
# the import of the target module to fail.
|
|
|
|
|
|
def notify_module_loaded(module):
|
|
"""
|
|
Notify that a `module` has been loaded and invoke any post import hooks
|
|
registered against the module. If the module is not registered, this
|
|
function does nothing.
|
|
"""
|
|
|
|
name = getattr(module, "__name__", None)
|
|
|
|
with _post_import_hooks_lock:
|
|
hooks = _post_import_hooks.pop(name, ())
|
|
|
|
# Note that the hook is called outside of the lock to avoid deadlocks if
|
|
# code run as a consequence of calling the module import hook in turn
|
|
# triggers a separate thread which tries to register an import hook.
|
|
|
|
for hook in hooks:
|
|
hook(module)
|
|
|
|
|
|
# A custom module import finder. This intercepts attempts to import
|
|
# modules and watches out for attempts to import target modules of
|
|
# interest. When a module of interest is imported, then any post import
|
|
# hooks which are registered will be invoked.
|
|
|
|
|
|
class _ImportHookLoader:
|
|
|
|
def load_module(self, fullname):
|
|
module = sys.modules[fullname]
|
|
notify_module_loaded(module)
|
|
|
|
return module
|
|
|
|
|
|
class _ImportHookChainedLoader(BaseObjectProxy):
|
|
|
|
def __init__(self, loader):
|
|
super(_ImportHookChainedLoader, self).__init__(loader)
|
|
|
|
if hasattr(loader, "load_module"):
|
|
self.__self_setattr__("load_module", self._self_load_module)
|
|
if hasattr(loader, "create_module"):
|
|
self.__self_setattr__("create_module", self._self_create_module)
|
|
if hasattr(loader, "exec_module"):
|
|
self.__self_setattr__("exec_module", self._self_exec_module)
|
|
|
|
def _self_set_loader(self, module):
|
|
# Set module's loader to self.__wrapped__ unless it's already set to
|
|
# something else. Import machinery will set it to spec.loader if it is
|
|
# None, so handle None as well. The module may not support attribute
|
|
# assignment, in which case we simply skip it. Note that we also deal
|
|
# with __loader__ not existing at all. This is to future proof things
|
|
# due to proposal to remove the attribute as described in the GitHub
|
|
# issue at https://github.com/python/cpython/issues/77458. Also prior
|
|
# to Python 3.3, the __loader__ attribute was only set if a custom
|
|
# module loader was used. It isn't clear whether the attribute still
|
|
# existed in that case or was set to None.
|
|
|
|
class UNDEFINED:
|
|
pass
|
|
|
|
if getattr(module, "__loader__", UNDEFINED) in (None, self):
|
|
try:
|
|
module.__loader__ = self.__wrapped__
|
|
except AttributeError:
|
|
pass
|
|
|
|
if (
|
|
getattr(module, "__spec__", None) is not None
|
|
and getattr(module.__spec__, "loader", None) is self
|
|
):
|
|
module.__spec__.loader = self.__wrapped__
|
|
|
|
def _self_load_module(self, fullname):
|
|
module = self.__wrapped__.load_module(fullname)
|
|
self._self_set_loader(module)
|
|
notify_module_loaded(module)
|
|
|
|
return module
|
|
|
|
# Python 3.4 introduced create_module() and exec_module() instead of
|
|
# load_module() alone. Splitting the two steps.
|
|
|
|
def _self_create_module(self, spec):
|
|
return self.__wrapped__.create_module(spec)
|
|
|
|
def _self_exec_module(self, module):
|
|
self._self_set_loader(module)
|
|
self.__wrapped__.exec_module(module)
|
|
notify_module_loaded(module)
|
|
|
|
|
|
class ImportHookFinder:
|
|
|
|
def __init__(self):
|
|
self.in_progress = {}
|
|
|
|
def find_module(self, fullname, path=None):
|
|
# If the module being imported is not one we have registered
|
|
# post import hooks for, we can return immediately. We will
|
|
# take no further part in the importing of this module.
|
|
|
|
with _post_import_hooks_lock:
|
|
if fullname not in _post_import_hooks:
|
|
return None
|
|
|
|
# When we are interested in a specific module, we will call back
|
|
# into the import system a second time to defer to the import
|
|
# finder that is supposed to handle the importing of the module.
|
|
# We set an in progress flag for the target module so that on
|
|
# the second time through we don't trigger another call back
|
|
# into the import system and cause a infinite loop.
|
|
|
|
if fullname in self.in_progress:
|
|
return None
|
|
|
|
self.in_progress[fullname] = True
|
|
|
|
# Now call back into the import system again.
|
|
|
|
try:
|
|
# For Python 3 we need to use find_spec().loader
|
|
# from the importlib.util module. It doesn't actually
|
|
# import the target module and only finds the
|
|
# loader. If a loader is found, we need to return
|
|
# our own loader which will then in turn call the
|
|
# real loader to import the module and invoke the
|
|
# post import hooks.
|
|
|
|
loader = getattr(find_spec(fullname), "loader", None)
|
|
|
|
if loader and not isinstance(loader, _ImportHookChainedLoader):
|
|
return _ImportHookChainedLoader(loader)
|
|
|
|
finally:
|
|
del self.in_progress[fullname]
|
|
|
|
def find_spec(self, fullname, path=None, target=None):
|
|
# Since Python 3.4, you are meant to implement find_spec() method
|
|
# instead of find_module() and since Python 3.10 you get deprecation
|
|
# warnings if you don't define find_spec().
|
|
|
|
# If the module being imported is not one we have registered
|
|
# post import hooks for, we can return immediately. We will
|
|
# take no further part in the importing of this module.
|
|
|
|
with _post_import_hooks_lock:
|
|
if fullname not in _post_import_hooks:
|
|
return None
|
|
|
|
# When we are interested in a specific module, we will call back
|
|
# into the import system a second time to defer to the import
|
|
# finder that is supposed to handle the importing of the module.
|
|
# We set an in progress flag for the target module so that on
|
|
# the second time through we don't trigger another call back
|
|
# into the import system and cause a infinite loop.
|
|
|
|
if fullname in self.in_progress:
|
|
return None
|
|
|
|
self.in_progress[fullname] = True
|
|
|
|
# Now call back into the import system again.
|
|
|
|
try:
|
|
# This should only be Python 3 so find_spec() should always
|
|
# exist so don't need to check.
|
|
|
|
spec = find_spec(fullname)
|
|
loader = getattr(spec, "loader", None)
|
|
|
|
if loader and not isinstance(loader, _ImportHookChainedLoader):
|
|
spec.loader = _ImportHookChainedLoader(loader)
|
|
|
|
return spec
|
|
|
|
finally:
|
|
del self.in_progress[fullname]
|
|
|
|
|
|
# Decorator for marking that a function should be called as a post
|
|
# import hook when the target module is imported.
|
|
|
|
|
|
def when_imported(name):
|
|
"""
|
|
Returns a decorator that registers the decorated function as a post import
|
|
hook for the module specified by `name`. The function will be called once
|
|
the module with the specified name is imported, and will be passed the
|
|
module as argument. If the module is already imported, the function will
|
|
be called immediately.
|
|
"""
|
|
|
|
def register(hook):
|
|
register_post_import_hook(hook, name)
|
|
return hook
|
|
|
|
return register
|