objgraph.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. """Find all objects reachable from a root object."""
  2. from __future__ import annotations
  3. import types
  4. import weakref
  5. from collections.abc import Iterable
  6. from typing import Final, Iterator, Mapping
  7. method_descriptor_type: Final = type(object.__dir__)
  8. method_wrapper_type: Final = type(object().__ne__)
  9. wrapper_descriptor_type: Final = type(object.__ne__)
  10. FUNCTION_TYPES: Final = (
  11. types.BuiltinFunctionType,
  12. types.FunctionType,
  13. types.MethodType,
  14. method_descriptor_type,
  15. wrapper_descriptor_type,
  16. method_wrapper_type,
  17. )
  18. ATTR_BLACKLIST: Final = {"__doc__", "__name__", "__class__", "__dict__"}
  19. # Instances of these types can't have references to other objects
  20. ATOMIC_TYPE_BLACKLIST: Final = {bool, int, float, str, type(None), object}
  21. # Don't look at most attributes of these types
  22. COLLECTION_TYPE_BLACKLIST: Final = {list, set, dict, tuple}
  23. # Don't return these objects
  24. TYPE_BLACKLIST: Final = {weakref.ReferenceType}
  25. def isproperty(o: object, attr: str) -> bool:
  26. return isinstance(getattr(type(o), attr, None), property)
  27. def get_edge_candidates(o: object) -> Iterator[tuple[object, object]]:
  28. # use getattr because mypyc expects dict, not mappingproxy
  29. if "__getattribute__" in getattr(type(o), "__dict__"): # noqa: B009
  30. return
  31. if type(o) not in COLLECTION_TYPE_BLACKLIST:
  32. for attr in dir(o):
  33. try:
  34. if attr not in ATTR_BLACKLIST and hasattr(o, attr) and not isproperty(o, attr):
  35. e = getattr(o, attr)
  36. if type(e) not in ATOMIC_TYPE_BLACKLIST:
  37. yield attr, e
  38. except AssertionError:
  39. pass
  40. if isinstance(o, Mapping):
  41. yield from o.items()
  42. elif isinstance(o, Iterable) and not isinstance(o, str):
  43. for i, e in enumerate(o):
  44. yield i, e
  45. def get_edges(o: object) -> Iterator[tuple[object, object]]:
  46. for s, e in get_edge_candidates(o):
  47. if isinstance(e, FUNCTION_TYPES):
  48. # We don't want to collect methods, but do want to collect values
  49. # in closures and self pointers to other objects
  50. if hasattr(e, "__closure__"):
  51. yield (s, "__closure__"), e.__closure__
  52. if hasattr(e, "__self__"):
  53. se = e.__self__
  54. if se is not o and se is not type(o) and hasattr(s, "__self__"):
  55. yield s.__self__, se
  56. else:
  57. if type(e) not in TYPE_BLACKLIST:
  58. yield s, e
  59. def get_reachable_graph(root: object) -> tuple[dict[int, object], dict[int, tuple[int, object]]]:
  60. parents = {}
  61. seen = {id(root): root}
  62. worklist = [root]
  63. while worklist:
  64. o = worklist.pop()
  65. for s, e in get_edges(o):
  66. if id(e) in seen:
  67. continue
  68. parents[id(e)] = (id(o), s)
  69. seen[id(e)] = e
  70. worklist.append(e)
  71. return seen, parents
  72. def get_path(
  73. o: object, seen: dict[int, object], parents: dict[int, tuple[int, object]]
  74. ) -> list[tuple[object, object]]:
  75. path = []
  76. while id(o) in parents:
  77. pid, attr = parents[id(o)]
  78. o = seen[pid]
  79. path.append((attr, o))
  80. path.reverse()
  81. return path