mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-07-30 05:32:30 +08:00
198 lines
6.5 KiB
Python
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
|