add get_object_graph

Author: matthiasdiener

pytools/debug.py

@@ -26,13 +26,14 @@
 .. autofunction:: make_unique_filesystem_object
 .. autofunction:: open_unique_debug_file
 .. autofunction:: refdebug
+.. autofunction:: get_object_graph
 .. autofunction:: get_object_cycles
 .. autofunction:: estimate_memory_usage
 
 """
 
 import sys
-from typing import Collection, List, Set
+from typing import Collection, Dict, List, Set, Tuple, Union
 
 from pytools import memoize
 
@@ -175,6 +176,48 @@ def is_excluded(o):
 
 # {{{ Find circular references
 
+from pytools.graph import GraphT
+
+
+def get_object_graph(objects: Collection[object],
+                     outside_objects: bool = False) -> GraphT:
+
+    """Create a graph out of *objects*, with graph edges representing
+    references between objects.
+
+    :arg objects: The objects to build the graph.
+    :arg outside_objects: Include objects not in *objects* in the graph when
+        ``True``.
+
+    :returns: A :class:`~pytools.graph.GraphT` with the references between objects.
+    """
+
+    import gc
+    res: Dict[object, List[object]] = {}
+
+    def hash_unhashable(obj: object) -> Union[object, Tuple[int, str]]:
+        try:
+            hash(obj)
+        except TypeError:
+            return (id(obj), str(obj))
+        else:
+            return obj
+
+    # Collect objects first to differentiate to outside objects later
+    for obj in objects:
+        res[hash_unhashable(obj)] = []
+
+    for obj in objects:
+        refs = gc.get_referents(obj)
+        obj = hash_unhashable(obj)
+        for r in refs:
+            r = hash_unhashable(r)
+            if r in res or outside_objects:
+                res.setdefault(r, []).append(obj)
+
+    return res
+
+
 # Based on https://code.activestate.com/recipes/523004-find-cyclical-references/
 
 def get_object_cycles(objects: Collection[object]) -> List[List[object]]:

test/test_debug.py

@@ -47,3 +47,36 @@ def test_get_object_cycles():
     b.append(a)
 
     assert len(get_object_cycles([a, b])) == 2
+
+
+def test_get_object_graph():
+    from pytools.debug import get_object_graph
+
+    assert get_object_graph([]) == {}
+
+    a = (1,)
+    b = (2,)
+    c = (a, b)
+    assert get_object_graph([a]) == {(1,): []}
+    assert get_object_graph([a, b]) == {(1,): [], (2,): []}
+    assert get_object_graph([a, b, c]) == {((1,), (2,)): [],      # c: []
+                                           (1,): [((1,), (2,))],  # a: [c]
+                                           (2,): [((1,), (2,))]}  # b: [c]
+
+    a = {}
+    b = {"a": a}
+    a["b"] = b
+
+    assert get_object_graph([a, b]) == {
+        (id(a), "{'b': {'a': {...}}}"): [(id(b), "{'a': {'b': {...}}}")],
+        (id(b), "{'a': {'b': {...}}}"): [(id(a), "{'b': {'a': {...}}}")]}
+
+    b = [42, 4]
+    a = [1, 2, 3, 4, 5, b]
+    b.append(a)
+
+    assert get_object_graph([a, b]) == {
+        (id(a), '[1, 2, 3, 4, 5, [42, 4, [...]]]'):
+            [(id(b), '[42, 4, [1, 2, 3, 4, 5, [...]]]')],
+        (id(b), '[42, 4, [1, 2, 3, 4, 5, [...]]]'):
+            [(id(a), '[1, 2, 3, 4, 5, [42, 4, [...]]]')]}