trace_dependencies.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. # mypy: allow-untyped-defs
  2. import sys
  3. from collections.abc import Callable, Iterable
  4. from typing import Any
  5. __all__ = ["trace_dependencies"]
  6. def trace_dependencies(
  7. callable: Callable[[Any], Any], inputs: Iterable[tuple[Any, ...]]
  8. ) -> list[str]:
  9. """Trace the execution of a callable in order to determine which modules it uses.
  10. Args:
  11. callable: The callable to execute and trace.
  12. inputs: The input to use during tracing. The modules used by 'callable' when invoked by each set of inputs
  13. are union-ed to determine all modules used by the callable for the purpooses of packaging.
  14. Returns: A list of the names of all modules used during callable execution.
  15. """
  16. modules_used = set()
  17. def record_used_modules(frame, event, arg):
  18. # If the event being profiled is not a Python function
  19. # call, there is nothing to do.
  20. if event != "call":
  21. return
  22. # This is the name of the function that was called.
  23. name = frame.f_code.co_name
  24. module = None
  25. # Try to determine the name of the module that the function
  26. # is in:
  27. # 1) Check the global namespace of the frame.
  28. # 2) Check the local namespace of the frame.
  29. # 3) To handle class instance method calls, check
  30. # the attribute named 'name' of the object
  31. # in the local namespace corresponding to "self".
  32. if name in frame.f_globals:
  33. module = frame.f_globals[name].__module__
  34. elif name in frame.f_locals:
  35. module = frame.f_locals[name].__module__
  36. elif "self" in frame.f_locals:
  37. method = getattr(frame.f_locals["self"], name, None)
  38. module = method.__module__ if method else None
  39. # If a module was found, add it to the set of used modules.
  40. if module:
  41. modules_used.add(module)
  42. try:
  43. # Attach record_used_modules as the profiler function.
  44. sys.setprofile(record_used_modules)
  45. # Execute the callable with all inputs.
  46. for inp in inputs:
  47. callable(*inp)
  48. finally:
  49. # Detach the profiler function.
  50. sys.setprofile(None)
  51. return list(modules_used)