Display graph in IPython.

objgraph.py

@@ -38,6 +38,12 @@
 import sys
 import itertools
 
+try:
+    # Python 2.x compatibility
+    from StringIO import StringIO
+except ImportError:
+    from io import StringIO
+
 try:
     from types import InstanceType
 except ImportError:
@@ -64,6 +70,14 @@
     # Python 3.x compatibility
     iteritems = dict.items
 
+try:
+    __IPYTHON__
+    import graphviz
+except (NameError, ImportError):
+    IS_INTERACTIVE = False
+else:
+    IS_INTERACTIVE = True
+
 
 def _isinstance(object, classinfo):
     """Return whether an object is an instance of a class or its subclass.
@@ -468,12 +482,12 @@ class name part.
     # module because you'll end up in sys.modules and explode the
     # graph with useless clutter.  That's why we're specifying
     # cull_func here, but not in show_graph().
-    _show_graph(objs, max_depth=max_depth, extra_ignore=extra_ignore,
-                filter=filter, too_many=too_many, highlight=highlight,
-                edge_func=gc.get_referrers, swap_source_target=False,
-                filename=filename, output=output, extra_info=extra_info,
-                refcounts=refcounts, shortnames=shortnames,
-                cull_func=is_proper_module)
+    return _show_graph(objs, max_depth=max_depth, extra_ignore=extra_ignore,
+                       filter=filter, too_many=too_many, highlight=highlight,
+                       edge_func=gc.get_referrers, swap_source_target=False,
+                       filename=filename, output=output, extra_info=extra_info,
+                       refcounts=refcounts, shortnames=shortnames,
+                       cull_func=is_proper_module)
 
 
 def show_refs(objs, max_depth=3, extra_ignore=(), filter=None, too_many=10,
@@ -536,11 +550,12 @@ def show_refs(objs, max_depth=3, extra_ignore=(), filter=None, too_many=10,
     .. versionchanged:: 2.0
        New parameter: ``output``.
     """
-    _show_graph(objs, max_depth=max_depth, extra_ignore=extra_ignore,
-                filter=filter, too_many=too_many, highlight=highlight,
-                edge_func=gc.get_referents, swap_source_target=True,
-                filename=filename, extra_info=extra_info, refcounts=refcounts,
-                shortnames=shortnames, output=output)
+    return _show_graph(objs, max_depth=max_depth, extra_ignore=extra_ignore,
+                       filter=filter, too_many=too_many, highlight=highlight,
+                       edge_func=gc.get_referents, swap_source_target=True,
+                       filename=filename, extra_info=extra_info,
+                       refcounts=refcounts, shortnames=shortnames,
+                       output=output)
 
 
 def show_chain(*chains, **kw):
@@ -654,6 +669,8 @@ def _show_graph(objs, edge_func, swap_source_target,
                 cull_func=None):
     if not _isinstance(objs, (list, tuple)):
         objs = [objs]
+
+    is_interactive = False
     if filename and output:
         raise ValueError('Cannot specify both output and filename.')
     elif output:
@@ -662,14 +679,18 @@ def _show_graph(objs, edge_func, swap_source_target,
         f = codecs.open(filename, 'w', encoding='utf-8')
         dot_filename = filename
     else:
-        fd, dot_filename = tempfile.mkstemp(prefix='objgraph-',
-                                            suffix='.dot', text=True)
-        f = os.fdopen(fd, "w")
-        if getattr(f, 'encoding', None):
-            # Python 3 will wrap the file in the user's preferred encoding
-            # Re-wrap it for utf-8
-            import io
-            f = io.TextIOWrapper(f.detach(), 'utf-8')
+        if IS_INTERACTIVE:
+            is_interactive = True
+            f = StringIO()
+        else:
+            fd, dot_filename = tempfile.mkstemp(prefix='objgraph-',
+                                                suffix='.dot', text=True)
+            f = os.fdopen(fd, "w")
+            if getattr(f, 'encoding', None):
+                # Python 3 will wrap the file in the user's preferred encoding
+                # Re-wrap it for utf-8
+                import io
+                f = io.TextIOWrapper(f.detach(), 'utf-8')
     f.write('digraph ObjectGraph {\n'
             '  node[shape=box, style=filled, fillcolor=white];\n')
     queue = []
@@ -767,13 +788,18 @@ def _show_graph(objs, edge_func, swap_source_target,
             f.write('  too_many_%s[fontcolor=white];\n'
                     % (_obj_node_id(target)))
     f.write("}\n")
+
     if output:
         return
-    # The file should only be closed if this function was in charge of opening
-    # the file.
-    f.close()
-    print("Graph written to %s (%d nodes)" % (dot_filename, nodes))
-    _present_graph(dot_filename, filename)
+
+    if is_interactive:
+        return graphviz.Source(f.getvalue())
+    else:
+        # The file should only be closed if this function was in charge of
+        # opening the file.
+        f.close()
+        print("Graph written to %s (%d nodes)" % (dot_filename, nodes))
+        _present_graph(dot_filename, filename)
 
 
 def _present_graph(dot_filename, filename=None):