Files
2020-04-08 10:39:44 -07:00

198 lines
6.5 KiB
Python

"""
This module is based off of wrapt.importer (wrapt==1.11.0)
https://github.com/GrahamDumpleton/wrapt/blob/4bcd190457c89e993ffcfec6dad9e9969c033e9e/src/wrapt/importer.py#L127-L136
The reasoning for this is that wrapt.importer does not provide a mechanism to
remove the import hooks and that wrapt removes the hooks after they are fired.
So this module differs from wrapt.importer in that:
- removes unnecessary functionality (like allowing hooks to be import paths)
- deregister_post_import_hook is introduced to remove hooks
- the values of _post_import_hooks can only be lists (instead of allowing None)
- notify_module_loaded is modified to not remove the hooks when they are
fired.
"""
import sys
import threading
from ..compat import PY3
from ..internal.logger import get_logger
from ..utils import get_module_name
from ..vendor.wrapt.decorators import synchronized
log = get_logger(__name__)
_post_import_hooks = {}
_post_import_hooks_init = False
_post_import_hooks_lock = threading.RLock()
@synchronized(_post_import_hooks_lock)
def register_post_import_hook(name, hook):
"""
Registers a module import hook, ``hook`` for a module with name ``name``.
If the module is already imported the hook is called immediately and a
debug message is logged since this should not be expected in our use-case.
:param name: Name of the module (full dotted path)
:type name: str
:param hook: Callable to be invoked with the module when it is imported.
:type hook: Callable
:return:
"""
# 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())
hooks = _post_import_hooks.get(name, [])
if hook in hooks:
log.debug('hook "%s" already exists on module "%s"', hook, name)
return
module = sys.modules.get(name, None)
# If the module has been imported already fire the hook and log a debug msg.
if module:
log.debug('module "%s" already imported, firing hook', name)
hook(module)
hooks.append(hook)
_post_import_hooks[name] = hooks
@synchronized(_post_import_hooks_lock)
def notify_module_loaded(module):
"""
Indicate that a module has been loaded. Any post import hooks which were
registered for the target module will be invoked.
Any raised exceptions will be caught and an error message indicating that
the hook failed.
:param module: The module being loaded
:type module: ``types.ModuleType``
"""
name = get_module_name(module)
hooks = _post_import_hooks.get(name, [])
for hook in hooks:
try:
hook(module)
except Exception:
log.warning('hook "%s" for module "%s" failed', hook, name, exc_info=True)
class _ImportHookLoader(object):
"""
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.
"""
def load_module(self, fullname):
module = sys.modules[fullname]
notify_module_loaded(module)
return module
class _ImportHookChainedLoader(object):
def __init__(self, loader):
self.loader = loader
def load_module(self, fullname):
module = self.loader.load_module(fullname)
notify_module_loaded(module)
return module
class ImportHookFinder:
def __init__(self):
self.in_progress = {}
@synchronized(_post_import_hooks_lock)
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.
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:
if PY3:
# 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
loader = importlib.util.find_spec(fullname).loader
except (ImportError, AttributeError):
loader = importlib.find_loader(fullname, path)
if loader:
return _ImportHookChainedLoader(loader)
else:
# For Python 2 we don't have much choice but to
# call back in to __import__(). This will
# actually cause the module to be imported. If no
# module could be found then ImportError will be
# raised. Otherwise we return a loader which
# returns the already loaded module and invokes
# the post import hooks.
__import__(fullname)
return _ImportHookLoader()
finally:
del self.in_progress[fullname]
@synchronized(_post_import_hooks_lock)
def deregister_post_import_hook(modulename, hook):
"""
Deregisters post import hooks for a module given the module name and a hook
that was previously installed.
:param modulename: Name of the module the hook is installed on.
:type: str
:param hook: The hook to remove (the function itself)
:type hook: Callable
:return: whether a hook was removed or not
"""
if modulename not in _post_import_hooks:
return False
hooks = _post_import_hooks[modulename]
try:
hooks.remove(hook)
return True
except ValueError:
return False