docs/sphinxext/plot_directive.py

"""A special directive for including a matplotlib plot.

Given a path to a .py file, it includes the source code inline, then:

- On HTML, will include a .png with a link to a high-res .png.

- On LaTeX, will include a .pdf

This directive supports all of the options of the `image` directive,
except for `target` (since plot will add its own target).

Additionally, if the :include-source: option is provided, the literal
source will be included inline, as well as a link to the source.
"""

import sys, os, glob, shutil
from docutils.parsers.rst import directives

try:
    # docutils 0.4
    from docutils.parsers.rst.directives.images import align
except ImportError:
    # docutils 0.5
    from docutils.parsers.rst.directives.images import Image
    align = Image.align


import matplotlib
import IPython.Shell
matplotlib.use('Agg')
import matplotlib.pyplot as plt

mplshell = IPython.Shell.MatplotlibShell('mpl')

options = {'alt': directives.unchanged,
           'height': directives.length_or_unitless,
           'width': directives.length_or_percentage_or_unitless,
           'scale': directives.nonnegative_int,
           'align': align,
           'class': directives.class_option,
           'include-source': directives.flag }

template = """
.. htmlonly::

   [`source code <../%(srcdir)s/%(basename)s.py>`__,
   `png <../%(srcdir)s/%(basename)s.hires.png>`__,
   `pdf <../%(srcdir)s/%(basename)s.pdf>`__]

   .. image:: ../%(srcdir)s/%(basename)s.png
%(options)s

.. latexonly::
   .. image:: ../%(srcdir)s/%(basename)s.pdf
%(options)s

"""

def makefig(fullpath, outdir):
    """
    run a pyplot script and save the low and high res PNGs and a PDF in _static
    """

    fullpath = str(fullpath)  # todo, why is unicode breaking this
    formats = [('png', 100),
               ('hires.png', 200),
               ('pdf', 72),
               ]

    basedir, fname = os.path.split(fullpath)
    basename, ext = os.path.splitext(fname)
    all_exists = True

    if basedir != outdir:
        shutil.copyfile(fullpath, os.path.join(outdir, fname))

    for format, dpi in formats:
        outname = os.path.join(outdir, '%s.%s' % (basename, format))
        if not os.path.exists(outname):
            all_exists = False
            break

    if all_exists:
        print '    already have %s'%fullpath
        return

    print '    building %s'%fullpath
    plt.close('all')    # we need to clear between runs
    matplotlib.rcdefaults()

    mplshell.magic_run(fullpath)
    for format, dpi in formats:
        outname = os.path.join(outdir, '%s.%s' % (basename, format))
        if os.path.exists(outname): continue
        plt.savefig(outname, dpi=dpi)

def run(arguments, options, state_machine, lineno):
    reference = directives.uri(arguments[0])
    basedir, fname = os.path.split(reference)
    basename, ext = os.path.splitext(fname)

    # todo - should we be using the _static dir for the outdir, I am
    # not sure we want to corrupt that dir with autogenerated files
    # since it also has permanent files in it which makes it difficult
    # to clean (save an rm -rf followed by and svn up)
    srcdir = 'pyplots'

    makefig(os.path.join(srcdir, reference), srcdir)

    # todo: it is not great design to assume the makefile is putting
    # the figs into the right place, so we may want to do that here instead.

    if options.has_key('include-source'):
        lines = ['.. literalinclude:: ../pyplots/%(reference)s' % locals()]
        del options['include-source']
    else:
        lines = []

    options = ['      :%s: %s' % (key, val) for key, val in
               options.items()]
    options = "\n".join(options)

    lines.extend((template % locals()).split('\n'))

    state_machine.insert_input(
        lines, state_machine.input_lines.source(0))
    return []


try:
    from docutils.parsers.rst import Directive
except ImportError:
    from docutils.parsers.rst.directives import _directives

    def plot_directive(name, arguments, options, content, lineno,
                       content_offset, block_text, state, state_machine):
        return run(arguments, options, state_machine, lineno)
    plot_directive.__doc__ = __doc__
    plot_directive.arguments = (1, 0, 1)
    plot_directive.options = options

    _directives['plot'] = plot_directive
else:
    class plot_directive(Directive):
        required_arguments = 1
        optional_arguments = 0
        final_argument_whitespace = True
        option_spec = options
        def run(self):
            return run(self.arguments, self.options,
                       self.state_machine, self.lineno)
    plot_directive.__doc__ = __doc__

    directives.register_directive('plot', plot_directive)