chore(internal): better support for Python legacy versions in ModuleWatchdog

ddtrace/internal/module.py

@@ -122,26 +122,24 @@ def _resolve(path):
 # https://github.com/GrahamDumpleton/wrapt/blob/df0e62c2740143cceb6cafea4c306dae1c559ef8/src/wrapt/importer.py
 
 if PY2:
+    import pkgutil
+
     find_spec = ModuleSpec = None
     Loader = object
+
+    find_loader = pkgutil.find_loader
+
 else:
     from importlib.abc import Loader
     from importlib.machinery import ModuleSpec
     from importlib.util import find_spec
 
+    def find_loader(fullname):
+        # type: (str) -> Optional[Loader]
+        return getattr(find_spec(fullname), "loader", None)
 
-# DEV: This is used by Python 2 only
-class _ImportHookLoader(object):
-    def __init__(self, callback):
-        # type: (Callable[[ModuleType], None]) -> None
-        self.callback = callback
 
-    def load_module(self, fullname):
-        # type: (str) -> ModuleType
-        module = sys.modules[fullname]
-        self.callback(module)
-
-        return module
+LEGACY_DICT_COPY = sys.version_info < (3, 6)
 
 
 class _ImportHookChainedLoader(Loader):
@@ -283,6 +281,18 @@ def __delitem__(self, name):
 
     def __getattribute__(self, name):
         # type: (str) -> Any
+        if LEGACY_DICT_COPY and name == "keys":
+            # This is a potential attempt to make a copy of sys.modules using
+            # dict(sys.modules) on a Python version that uses the C API to
+            # perform the operation. Since we are an instance of a dict, this
+            # means that we will end up looking like the empty dict, so we take
+            # this chance to actually look like sys.modules.
+            # NOTE: This is a potential source of memory leaks. However, we
+            # expect this to occur only on defunct Python versions, and only
+            # during special code executions, like test runs.
+            super(ModuleWatchdog, self).clear()
+            super(ModuleWatchdog, self).update(self._modules)
+
         try:
             return super(ModuleWatchdog, self).__getattribute__("_modules").__getattribute__(name)
         except AttributeError:
@@ -308,16 +318,20 @@ def find_module(self, fullname, path=None):
         self._finding.add(fullname)
 
         try:
-            if PY2:
-                __import__(fullname)
-                return _ImportHookLoader(self.after_import)
-
-            loader = getattr(find_spec(fullname), "loader", None)
+            loader = find_loader(fullname)
             if loader is not None:
                 if not isinstance(loader, _ImportHookChainedLoader):
                     loader = _ImportHookChainedLoader(loader)
 
-                loader.add_callback(type(self), self.after_import)
+                if PY2:
+                    # With Python 2 we don't get all the finders invoked, so we
+                    # make sure we register all the callbacks at the earliest
+                    # opportunity.
+                    for finder in sys.meta_path:
+                        if isinstance(finder, ModuleWatchdog):
+                            loader.add_callback(type(finder), finder.after_import)
+                else:
+                    loader.add_callback(type(self), self.after_import)
 
                 return loader
 

tests/internal/test_module.py

@@ -310,3 +310,20 @@ class Bob(BaseCollector):
 
     Bob.uninstall()
     Alice.uninstall()
+
+
+def test_module_watchdog_dict_shallow_copy():
+    # Save original reference to sys.modules
+    original_sys_modules = sys.modules
+
+    ModuleWatchdog.install()
+
+    # Ensure that we have replaced sys.modules
+    assert original_sys_modules is not sys.modules
+
+    # Make a shallow copy of both using the dict constructor
+    original_modules = set(dict(original_sys_modules).keys())
+    new_modules = set(dict(sys.modules).keys())
+
+    # Ensure that they match
+    assert original_modules == new_modules