hierarchymagic.py

"""
`%hierarchy`, `%%dot` and `%dotstr` magics for IPython
===========================================

This extension provides three magics.

First magic is ``%hierarchy``.  This magic command draws hierarchy of
given class or the class of given instance.  For example, the
following shows class hierarchy of currently running IPython shell.::

    %hierarchy get_ipython()


Second magic is ``%%dot``.  You can write graphiz dot language in a
cell using this magic.  Example::

    %%dot -- -Kfdp
    digraph G {
        a->b; b->c; c->d; d->b; d->a;
    }

Third magic is ``%%dotstr``. When you have a string written in the dot
language, use this function. E.g. when using pydot objects. Example::

    %dotstr -f svg -o-Kfdp pydot_obj.to_string()


License for ipython-hierarchymagic
----------------------------------

ipython-hierarchymagic is licensed under the term of the Simplified
BSD License (BSD 2-clause license), as follows:

Copyright (c) 2012 Takafumi Arakaki
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.

Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


License for Sphinx
------------------

`run_dot` function and `HierarchyMagic._class_name` method in this
extension heavily based on Sphinx code `sphinx.ext.graphviz.render_dot`
and `InheritanceGraph.class_name`.

Copyright notice for Sphinx can be found below.

Copyright (c) 2007-2011 by the Sphinx team (see AUTHORS file).
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

* Redistributions of source code must retain the above copyright
  notice, this list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright
  notice, this list of conditions and the following disclaimer in the
  documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""

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 sphinx.ext.inheritance_diagram import InheritanceGraph


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.encode('utf-8'))
    except (OSError, IOError) as err:
        if err.errno != EPIPE:
            raise
        wentwrong = True
    except IOError as 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.decode('utf-8')))
    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 dot(self, line, cell):
        """ Draw a figure from the cell contents (using Graphviz dot)."""
        args = parse_argstring(self.dot, line)
        self._dot(args, cell)

    @magic_arguments()
    @argument(
        '-f', '--format', default='png', choices=('png', 'svg'),
        help='Output format (png/svg).'
    )
    @argument(
        '-o', '--options', default=[], nargs='*',
        help='Options passed to the `dot` command.'
    )
    @argument(
        'strings', nargs='*',
        help='Strings in graphviz syntax.'
    )
    @line_magic
    def dotstr(self, line):
        """ Draw a figure from a string (using Graphviz dot). """
        args = parse_argstring(self.dotstr, line)
        for s in args.strings:
            self._dot(args, self.shell.ev(s))

    def _dot(self, args, string):
        """ Run the Graphviz dot command. """
        # Parse the arguments
        image = run_dot(string, args.options, format=args.format)
        if args.format == 'png':
            display_png(image, raw=True)
        elif args.format == 'svg':
            display_svg(image, raw=True)


class FoldedInheritanceGraph(InheritanceGraph):

    def __init__(self, *args, **kwds):
        self._width = kwds.pop('width', 40)
        super(FoldedInheritanceGraph, self).__init__(*args, **kwds)

    @staticmethod
    def _foldclassname(classname, width):
        r"""
        Split `classname` in newlines if the width is wider than `width`.

        >>> fold = FoldedInheritanceGraph._foldclassname
        >>> fold('aaa.bbb.ccc', 7)
        'aaa.bbb\\n.ccc'
        >>> fold('aaa.bbb.ccc', 3)
        'aaa\\n.bbb\\n.ccc'
        >>> identity = lambda x, y: ''.join(fold(x, y).split('\\n'))
        >>> identity('aaa.bbb.ccc', 7)
        'aaa.bbb.ccc'
        >>> identity('aaa.bbb.ccc', 3)
        'aaa.bbb.ccc'

        """
        parts = classname.split('.')
        lines = []
        chunk = [parts.pop(0)]
        for p in parts:
            if len('.'.join(chunk + [p])) > width:
                lines.append('.'.join(chunk))
                chunk = [p]
            else:
                chunk.append(p)
        lines.append('.'.join(chunk))
        return '\\n.'.join(lines)

    def _class_info(self, *args, **kwds):
        class_info = super(FoldedInheritanceGraph, self) \
            ._class_info(*args, **kwds)
        width = self._width

        def fold(elem):
            (nodename, fullname, baselist, tooltip) = elem
            nodename = self._foldclassname(nodename, width)
            baselist = [self._foldclassname(b, width) for b in baselist]
            return (nodename, fullname, baselist, tooltip)

        return map(fold, class_info)


@magics_class
class HierarchyMagic(Magics):

    @magic_arguments()
    @argument(
        '-r', '--rankdir', default='TB',
        help='direction of the hierarchy graph (default: %(default)s)'
    )
    @argument(
        '-s', '--size', default='5.0, 12.0',
        help='size of the generated figure (default: %(default)s)',
    )
    @argument(
        '-w', '--name-width', default=80, type=int,
        help='width of each nodes in character length (default: %(default)s)',
    )
    @argument(
        'object', nargs='+',
        help='Class hierarchy of these classes or objects will be drawn',
    )
    @line_magic
    def hierarchy(self, parameter_s=''):
        """Draw hierarchy of a given class."""
        args = parse_argstring(self.hierarchy, parameter_s)
        objects = map(self.shell.ev, args.object)
        clslist = map(self._object_to_class, objects)
        namelist = map(self._class_name, clslist)
        ig = FoldedInheritanceGraph(
            namelist, '',
            width=args.name_width)
        code = ig.generate_dot('inheritance_graph',
                               graph_attrs={'rankdir': args.rankdir,
                                            'size': '"{0}"'.format(args.size)})
        stdout = run_dot(code, format='png')
        display_png(stdout, raw=True)

    @staticmethod
    def _object_to_class(obj):
        if isinstance(obj, type):
            return obj
        elif hasattr(obj, "__class__"):
            return obj.__class__
        else:
            raise ValueError(
                "Given object {0} is not a class or an instance".format(obj))

    @staticmethod
    def _class_name(cls, parts=0):
        """Given a class object, return a fully-qualified name.

        This works for things I've tested in matplotlib so far, but may not be
        completely general.
        """
        module = cls.__module__
        if module == '__builtin__':
            fullname = cls.__name__
        else:
            fullname = '%s.%s' % (module, cls.__name__)
        if parts == 0:
            return fullname
        name_parts = fullname.split('.')
        return '.'.join(name_parts[-parts:])


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

_loaded = False