Move DD code into its own directory (#6)

reference/ddtrace/utils/hook.py

@@ -0,0 +1,197 @@
+"""
+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