Files
2026-07-13 13:22:34 +08:00

358 lines
13 KiB
Python

"""
NOTE: The contents of this file have been inlined from the wrapt package's source code
https://github.com/GrahamDumpleton/wrapt/blob/1.12.1/src/wrapt/importer.py.
Some modifications, have been made in order to:
- avoid duplicate registration of import hooks
- inline functions from dependent wrapt submodules rather than importing them.
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.
It also extends the functionality to support custom hooks for import errors
(as opposed to only successful imports).
"""
import importlib.resources
import sys
import threading
string_types = (str,)
# from .decorators import synchronized
# NOTE: Instead of using this import (from wrapt's decorator module, see
# https://github.com/GrahamDumpleton/wrapt/blob/68316bea668fd905a4acb21f37f12596d8c30d80/src/wrapt/decorators.py#L430-L456),
# we define a decorator with similar behavior that acquires a lock while calling the decorated
# function
def synchronized(lock):
def decorator(f):
# See e.g. https://www.python.org/dev/peps/pep-0318/#examples
def new_fn(*args, **kwargs):
with lock:
return f(*args, **kwargs)
return new_fn
return decorator
# 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 = {}
_post_import_hooks_lock = threading.RLock()
# A dictionary for any import hook error handlers to be triggered when the
# target module import fails.
_import_error_hooks = {}
_import_error_hooks_lock = threading.RLock()
_import_hook_finder_init = False
# 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_generic_import_hook(hook, name, hook_dict, overwrite):
# Create a deferred import hook if hook is a string name rather than
# a callable function.
if isinstance(hook, string_types):
hook = _create_import_hook_from_string(hook)
# Automatically install the import hook finder if it has not already
# been installed.
global _import_hook_finder_init
if not _import_hook_finder_init:
_import_hook_finder_init = True
sys.meta_path.insert(0, ImportHookFinder())
# Determine if any prior registration of an import hook for
# the target modules has occurred and act appropriately.
hooks = hook_dict.get(name, None)
if hooks is None:
# No prior registration of import hooks for the target
# module. We need to check whether the module has already been
# imported. If it has we fire the hook immediately and add an
# empty list to the registry to indicate that the module has
# already been imported and hooks have fired. Otherwise add
# the post import hook to the registry.
module = sys.modules.get(name, None)
if module is not None:
hook_dict[name] = []
hook(module)
else:
hook_dict[name] = [hook]
elif hooks == []:
# A prior registration of import hooks for the target
# module was done and the hooks already fired. Fire the hook
# immediately.
module = sys.modules[name]
hook(module)
else:
# A prior registration of import hooks for the target
# module was done but the module has not yet been imported.
def hooks_equal(existing_hook, hook):
if hasattr(existing_hook, "__name__") and hasattr(hook, "__name__"):
return existing_hook.__name__ == hook.__name__
else:
return False
if overwrite:
hook_dict[name] = [
existing_hook
for existing_hook in hook_dict[name]
if not hooks_equal(existing_hook, hook)
]
hook_dict[name].append(hook)
@synchronized(_import_error_hooks_lock)
def register_import_error_hook(hook, name, overwrite=True):
"""
Args:
hook: A function or string entrypoint to invoke when the specified module is imported
and an error occurs.
name: The name of the module for which to fire the hook at import error detection time.
overwrite: Specifies the desired behavior when a preexisting hook for the same
function / entrypoint already exists for the specified module. If `True`,
all preexisting hooks matching the specified function / entrypoint will be
removed and replaced with a single instance of the specified `hook`.
"""
register_generic_import_hook(hook, name, _import_error_hooks, overwrite)
@synchronized(_post_import_hooks_lock)
def register_post_import_hook(hook, name, overwrite=True):
"""
Args:
hook: A function or string entrypoint to invoke when the specified module is imported.
name: The name of the module for which to fire the hook at import time.
overwrite: Specifies the desired behavior when a preexisting hook for the same
function / entrypoint already exists for the specified module. If `True`,
all preexisting hooks matching the specified function / entrypoint will be
removed and replaced with a single instance of the specified `hook`.
"""
register_generic_import_hook(hook, name, _post_import_hooks, overwrite)
@synchronized(_post_import_hooks_lock)
def get_post_import_hooks(name):
return _post_import_hooks.get(name)
# Register post import hooks defined as package entry points.
def _create_import_hook_from_entrypoint(entrypoint):
def import_hook(module):
__import__(entrypoint.module_name)
callback = sys.modules[entrypoint.module_name]
for attr in entrypoint.attrs:
callback = getattr(callback, attr)
return callback(module)
return import_hook
def discover_post_import_hooks(group):
for entrypoint in (
resource.name
for resource in importlib.resources.files(group).iterdir()
if resource.is_file()
):
callback = _create_import_hook_from_entrypoint(entrypoint)
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):
# Snapshot hooks under the lock, then release it before firing them. Hooks
# may trigger imports that re-enter the import machinery on another thread
# and try to acquire this lock, which would deadlock against per-module
# import locks held by that thread.
name = getattr(module, "__name__", None)
with _post_import_hooks_lock:
hooks = _post_import_hooks.get(name) or []
if hooks:
_post_import_hooks[name] = []
for hook in hooks:
hook(module)
def notify_module_import_error(module_name):
with _import_error_hooks_lock:
# Error hooks differ from post import hooks, in that we don't clear the
# hook as soon as it fires.
hooks = list(_import_error_hooks.get(module_name) or [])
for hook in hooks:
hook(module_name)
# 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 _ImportHookChainedLoader:
def __init__(self, loader):
self.loader = loader
def load_module(self, fullname):
try:
module = self.loader.load_module(fullname)
notify_module_loaded(module)
except (ImportError, AttributeError):
notify_module_import_error(fullname)
raise
return module
class ImportHookFinder:
def __init__(self):
self.in_progress = {}
@synchronized(_post_import_hooks_lock)
@synchronized(_import_error_hooks_lock)
def find_module(self, fullname, path=None):
# If the module being imported is not one we have registered
# import hooks for, we can return immediately. We will
# take no further part in the importing of this module.
if fullname not in _post_import_hooks and fullname not in _import_error_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.
try:
import importlib.util # clint: disable=lazy-import
loader = importlib.util.find_spec(fullname).loader
# If an ImportError (or AttributeError) is encountered while finding the module,
# notify the hooks for import errors
except (ImportError, AttributeError):
notify_module_import_error(fullname)
loader = importlib.find_loader(fullname, path)
if loader:
return _ImportHookChainedLoader(loader)
finally:
del self.in_progress[fullname]
@synchronized(_post_import_hooks_lock)
@synchronized(_import_error_hooks_lock)
def find_spec(self, fullname, path, target=None):
# If the module being imported is not one we have registered
# import hooks for, we can return immediately. We will
# take no further part in the importing of this module.
if fullname not in _post_import_hooks and fullname not in _import_error_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:
import importlib.util # clint: disable=lazy-import
spec = importlib.util.find_spec(fullname)
# Replace the module spec's loader with a wrapped version that executes import
# hooks when the module is loaded
spec.loader = _ImportHookChainedLoader(spec.loader)
return spec
except (ImportError, AttributeError):
notify_module_import_error(fullname)
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.
# If error_handler is True, then apply the marked function as an import hook
# for import errors (instead of successful imports).
# It is assumed that all error hooks are added during driver start-up,
# and thus added prior to any import calls. If an error hook is added
# after a module has already failed the import, there's no guarantee
# that the hook will fire.
def when_imported(name, error_handler=False):
def register(hook):
if error_handler:
register_import_error_hook(hook, name)
else:
register_post_import_hook(hook, name)
return hook
return register