ipython_object_graphs.py

"""
An experimental extension for an IPython magic command to show object graphs.

Strongly based on tkf's code for the extension itself
and mine for turning a Python namespace into a GraphViz DOT source:

- https://github.com/tkf/ipython-hierarchymagic
- http://pypi.python.org/pypi/pyrels/0.1.1

Read respective copyrights there!

Dinu Gherman
2013-02-12
"""


from IPython.core.magic \
    import Magics, magics_class, line_magic, cell_magic
from IPython.core.magic_arguments \
    import argument, magic_arguments, parse_argstring
from IPython.core.display import display_png, display_svg

from pyrels.pyrels2dot import namespace2dot_str


def run_dot(code, options=[], format='png'):
    # mostly copied from sphinx.ext.graphviz.render_dot
    import os
    from subprocess import Popen, PIPE
    from sphinx.util.osutil import EPIPE, EINVAL

    dot_args = ['dot'] + options + ['-T', format]
    if os.name == 'nt':
        # Avoid opening shell window.
        # * https://github.com/tkf/ipython-hierarchymagic/issues/1
        # * http://stackoverflow.com/a/2935727/727827
        p = Popen(dot_args, stdout=PIPE, stdin=PIPE, stderr=PIPE,
                  creationflags=0x08000000)
    else:
        p = Popen(dot_args, stdout=PIPE, stdin=PIPE, stderr=PIPE)
    wentwrong = False
    try:
        # Graphviz may close standard input when an error occurs,
        # resulting in a broken pipe on communicate()
        stdout, stderr = p.communicate(code)
    except (OSError, IOError), err:
        if err.errno != EPIPE:
            raise
        wentwrong = True
    except IOError, err:
        if err.errno != EINVAL:
            raise
        wentwrong = True
    if wentwrong:
        # in this case, read the standard output and standard error streams
        # directly, to get the error message(s)
        stdout, stderr = p.stdout.read(), p.stderr.read()
        p.wait()
    if p.returncode != 0:
        raise RuntimeError('dot exited with error:\n[stderr]\n{0}'
                           .format(stderr))
    return stdout


@magics_class
class GraphvizMagic(Magics):

    @magic_arguments()
    @argument(
        '-f', '--format', default='png', choices=('png', 'svg'),
        help='output format (png/svg)'
    )
    @argument(
        'options', default=[], nargs='*',
        help='options passed to the `dot` command'
    )
    @cell_magic
    def objgraph(self, line, cell):
        """Draw a figure using Graphviz dot command."""
        
        args = parse_argstring(self.objgraph, line)

        dot = namespace2dot_str(eval(cell, _ip_globals))
        image = run_dot(dot, args.options, format=args.format)

        if args.format == 'png':
            display_png(image, raw=True)
        elif args.format == 'svg':
            display_svg(image, raw=True)


def load_ipython_extension(ip):
    """Load the extension in IPython."""
    
    global _loaded
    global _ip_globals
    if not _loaded:
        ip.register_magics(GraphvizMagic)
        _ip_globals = ip.user_global_ns
        _loaded = True
    else:
        _ip_globals = ip.user_global_ns


_loaded = False