Merge pull request #30 from IBM/TrackImportStack-27

Track import stack 27

Author: gabe-l-hart

import_tracker/__main__.py

@@ -17,10 +17,11 @@
 # Standard
 from concurrent.futures import ThreadPoolExecutor
 from types import ModuleType
-from typing import List, Optional, Set
+from typing import Dict, List, Optional, Set, Union
 import argparse
 import cmath
 import importlib
+import inspect
 import json
 import logging
 import os
@@ -54,9 +55,10 @@ def _get_import_parent_path(mod) -> str:
     return parent_path
 
 
-def _get_non_std_modules(mod_names: Set[str]) -> Set[str]:
+def _get_non_std_modules(mod_names: Union[Set[str], Dict[str, List[dict]]]) -> Set[str]:
     """Take a snapshot of the non-standard modules currently imported"""
-    return {
+    # Determine the names from the list that are non-standard
+    non_std_mods = {
         mod_name.split(".")[0]
         for mod_name, mod in sys.modules.items()
         if mod_name in mod_names
@@ -66,6 +68,17 @@ def _get_non_std_modules(mod_names: Set[str]) -> Set[str]:
         and mod_name.split(".")[0] != THIS_PACKAGE
     }
 
+    # If this is a set, just return it directly
+    if isinstance(mod_names, set):
+        return non_std_mods
+
+    # If it's a dict, limit to the non standard names
+    return {
+        mod_name: mod_vals
+        for mod_name, mod_vals in mod_names.items()
+        if mod_name in non_std_mods
+    }
+
 
 class _DeferredModule(ModuleType):
     """A _DeferredModule is a module subclass that wraps another module but imports
@@ -197,17 +210,16 @@ def exec_module(self, *_, **__):
 class ImportTrackerMetaFinder(importlib.abc.MetaPathFinder):
     """The ImportTrackerMetaFinder is a meta finder that is intended to be used
     at the front of the sys.meta_path to automatically track the imports for a
-    given library. It does this by looking at the call stack when a given import
-    is requested and tracking the upstream for each import made inside of the
-    target package.
-
-    NOTE: Since a stack trace is traversed on every import, this is very slow
-        and is intended only for a static build-time operation and should not be
-        used during the import phase of a library at runtime!
+    given library. It does this by deferring all imports which occur before the
+    target module has been seen, then collecting all imports seen until the
+    target import has completed.
     """
 
     def __init__(
-        self, tracked_module: str, side_effect_modules: Optional[List[str]] = None
+        self,
+        tracked_module: str,
+        side_effect_modules: Optional[List[str]] = None,
+        track_import_stack: bool = False,
     ):
         """Initialize with the name of the package being tracked
 
@@ -219,6 +231,12 @@ def __init__(
                 to perform required import tasks (e.g. global singleton
                 registries). These modules will be allowed to import regardless
                 of where they fall relative to the targeted module.
+            track_import_stack:  bool
+                If true, when imports are allowed through, their stack trace is
+                captured.
+                NOTE: This will cause a stack trace to be computed for every
+                    import in the tracked set, so it will be very slow and
+                    should only be used as a debugging tool on targeted imports.
         """
         self._tracked_module = tracked_module
         self._side_effect_modules = side_effect_modules or []
@@ -228,9 +246,14 @@ def __init__(
         log.debug2("Starting modules: %s", self._starting_modules)
         self._ending_modules = None
         self._deferred_modules = set()
+        self._track_import_stack = track_import_stack
+        self._import_stacks = {}
 
     def find_spec(
-        self, fullname: str, *args, **kwargs
+        self,
+        fullname: str,
+        *args,
+        **kwargs,
     ) -> Optional[importlib.machinery.ModuleSpec]:
         """The find_spec implementation for this finder tracks the source of the
         import call for the given module and determines if it is on the critical
@@ -249,6 +272,59 @@ def find_spec(
                 import is on the critical path, None will be returned to defer
                 to the rest of the "real" finders.
         """
+        # Do the main tracking logic
+        result = self._find_spec(fullname, *args, **kwargs)
+
+        # If this module is deferred, return it
+        if result is not None:
+            log.debug2("Returning deferred module for [%s]", fullname)
+            return result
+
+        # If this module is part of the set of modules belonging to the tracked
+        # module and stack tracing is enabled, grab all frames in the stack that
+        # come from the tracked module's package.
+        log.debug2(
+            "Stack tracking? %s, Ending modules set? %s",
+            self._track_import_stack,
+            self._ending_modules is not None,
+        )
+        if (
+            self._track_import_stack
+            and fullname != self._tracked_module
+            and not self._enabled
+        ):
+            stack = inspect.stack()
+            stack_info = []
+            for frame in stack:
+                frame_module_name = frame.frame.f_globals["__name__"].split(".")[0]
+                if frame_module_name == self._tracked_module_parts[0]:
+                    stack_info.append(
+                        {
+                            "filename": frame.filename,
+                            "lineno": frame.lineno,
+                            "code_context": [
+                                line.strip("\n") for line in frame.code_context
+                            ],
+                        }
+                    )
+
+            # NOTE: Under certain _strange_ cases, you can end up overwriting a
+            #   previous import stack here. I've only ever seen this happen with
+            #   pytest internals. Also, in this case the best we can do is just
+            #   keep the latest one.
+            log.debug2("Found %d stack frames for [%s]", len(stack_info), fullname)
+            self._import_stacks[fullname] = stack_info
+
+        # Let the module pass through
+        return None
+
+    def _find_spec(
+        self, fullname: str, *args, **kwargs
+    ) -> Optional[importlib.machinery.ModuleSpec]:
+        """This implements the core logic of find_spec. It is wrapped by the
+        public find_spec so that when an import is allowed, the stack can be
+        optionally tracked.
+        """
 
         # If this module fullname is one of the modules with known side-effects,
         # let it fall through
@@ -309,11 +385,17 @@ def get_all_new_modules(self) -> Set[str]:
         assert self._starting_modules is not None, f"Target module never impoted!"
         if self._ending_modules is None:
             self._set_ending_modules()
-        return {
+        mod_names = {
             mod
             for mod in self._ending_modules - self._starting_modules
             if not self._is_parent_module(mod)
         }
+        if self._track_import_stack:
+            return {
+                mod_name: self._import_stacks.get(mod_name, [])
+                for mod_name in mod_names
+            }
+        return mod_names
 
     ## Implementation Details ##
 
@@ -417,6 +499,13 @@ def main():
         default=None,
         help="Modules with known import-time side effect which should always be allowed to import",
     )
+    parser.add_argument(
+        "--track_import_stack",
+        "-t",
+        action="store_true",
+        default=False,
+        help="Store the stack trace of imports belonging to the tracked module",
+    )
     args = parser.parse_args()
 
     # Validate sets of args
@@ -442,7 +531,11 @@ def main():
         full_module_name = f"{args.package}{args.name}"
 
     # Create the tracking meta finder
-    tracker_finder = ImportTrackerMetaFinder(full_module_name, args.side_effect_modules)
+    tracker_finder = ImportTrackerMetaFinder(
+        tracked_module=full_module_name,
+        side_effect_modules=args.side_effect_modules,
+        track_import_stack=args.track_import_stack,
+    )
     sys.meta_path = [tracker_finder] + sys.meta_path
 
     # Do the import
@@ -480,6 +573,14 @@ def main():
             ]
         log.debug("Recursing on: %s", recursive_internals)
 
+        # Set up the kwargs for recursing
+        recursive_kwargs = dict(
+            log_level=log_level,
+            recursive=False,
+            side_effect_modules=args.side_effect_modules,
+            track_import_stack=args.track_import_stack,
+        )
+
         # Create the thread pool to manage the subprocesses
         if args.num_jobs > 0:
             pool = ThreadPoolExecutor(max_workers=args.num_jobs)
@@ -489,9 +590,7 @@ def main():
                     pool.submit(
                         track_module,
                         module_name=internal_downstream,
-                        log_level=log_level,
-                        recursive=False,
-                        side_effect_modules=args.side_effect_modules,
+                        **recursive_kwargs,
                     )
                 )
 
@@ -507,12 +606,10 @@ def main():
                     )
                     downstream_mapping.update(
                         track_module(
-                            module_name=internal_downstream,
-                            log_level=log_level,
-                            recursive=False,
-                            side_effect_modules=args.side_effect_modules,
+                            module_name=internal_downstream, **recursive_kwargs
                         )
                     )
+
                 # This is useful for catching errors caused by unexpected corner
                 # cases. If it's triggered, it's a sign of a bug in the library,
                 # so we don't have ways to explicitly exercise this in tests.
@@ -527,13 +624,19 @@ def main():
     # Get all of the downstreams for the module in question, including internals
     log.debug("Downstream Mapping: %s", downstream_mapping)
 
+    # Set up the output dict depending on whether or not the stack info is being
+    # tracked
+    if args.track_import_stack:
+        output_dict = {
+            key: dict(sorted(val.items())) for key, val in downstream_mapping.items()
+        }
+    else:
+        output_dict = {
+            key: sorted(list(val)) for key, val in downstream_mapping.items()
+        }
+
     # Print out the json dump
-    print(
-        json.dumps(
-            {key: sorted(list(val)) for key, val in downstream_mapping.items()},
-            indent=args.indent,
-        ),
-    )
+    print(json.dumps(output_dict, indent=args.indent))
 
 
 if __name__ == "__main__":

import_tracker/import_tracker.py

@@ -28,6 +28,7 @@ def track_module(
     num_jobs: int = 0,
     side_effect_modules: Optional[List[str]] = None,
     submodules: Optional[List[str]] = None,
+    track_import_stack: bool = False,
 ) -> Dict[str, List[str]]:
     """This function executes the tracking of a single module by launching a
     subprocess to execute this module against the target module. The
@@ -54,6 +55,8 @@ def track_module(
             fall relative to the targeted module.
         submodules:  Optional[List[str]]
             List of sub-modules to recurse on (only used when recursive set)
+        track_import_stack:  bool
+            Store the stack trace of imports belonging to the tracked module
 
     Returns:
         import_mapping:  Dict[str, List[str]]
@@ -87,6 +90,8 @@ def track_module(
         cmd += " --side_effect_modules " + " ".join(side_effect_modules)
     if submodules:
         cmd += " --submodules " + " ".join(submodules)
+    if track_import_stack:
+        cmd += " --track_import_stack"
 
     # Launch the process
     proc = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE, env=env)

test/test_main.py

@@ -6,6 +6,7 @@
 from contextlib import contextmanager
 import json
 import logging
+import os
 import sys
 
 # Third Party
@@ -198,3 +199,67 @@ def test_error_submodules_without_recursive():
     ):
         with pytest.raises(ValueError):
             main()
+
+
+def test_import_stack_tracking(capsys):
+    """Make sure that tracking the import stack works as expected"""
+    with cli_args(
+        "--name",
+        "inter_mod_deps",
+        "--recursive",
+        "--track_import_stack",
+    ):
+        main()
+    captured = capsys.readouterr()
+    assert captured.out
+    parsed_out = json.loads(captured.out)
+
+    assert set(parsed_out.keys()) == {
+        "inter_mod_deps",
+        "inter_mod_deps.submod1",
+        "inter_mod_deps.submod2",
+        "inter_mod_deps.submod2.foo",
+        "inter_mod_deps.submod2.bar",
+        "inter_mod_deps.submod3",
+        "inter_mod_deps.submod4",
+        "inter_mod_deps.submod5",
+    }
+
+    # Check one of the stacks to make sure it's correct
+    test_lib_dir = os.path.realpath(
+        os.path.join(
+            os.path.dirname(__file__),
+            "sample_libs",
+            "inter_mod_deps",
+        )
+    )
+    assert parsed_out["inter_mod_deps.submod2"] == {
+        "alog": [
+            {
+                "filename": f"{test_lib_dir}/submod1/__init__.py",
+                "lineno": 6,
+                "code_context": ["import alog"],
+            },
+            {
+                "filename": f"{test_lib_dir}/__init__.py",
+                "lineno": 17,
+                "code_context": [
+                    "from . import submod1, submod2, submod3, submod4, submod5"
+                ],
+            },
+        ],
+        "yaml": [
+            {
+                "filename": f"{test_lib_dir}/submod2/__init__.py",
+                "lineno": 6,
+                "code_context": ["import yaml"],
+            },
+            {
+                "filename": f"{test_lib_dir}/__init__.py",
+                "lineno": 17,
+                "code_context": [
+                    "from . import submod1, submod2, submod3, submod4, submod5"
+                ],
+            },
+        ],
+    }