Refactor plot:: directive somewhat

1 files changed, 324 insertions, 167 deletions

diff --git a/doc/sphinxext/plot_directive.py b/doc/sphinxext/plot_directive.py
index 5fa24791d..8a82aac74 100644
--- a/doc/sphinxext/plot_directive.py
+++ b/doc/sphinxext/plot_directive.py
@@ -1,34 +1,102 @@
-# plot_directive.py from matplotlib.sf.net
-"""A special directive for including a matplotlib plot.
+"""
+A special directive for generating a matplotlib plot.
+
+.. warning::
 
-Given a path to a .py file, it includes the source code inline, then:
+   This is a hacked version of plot_directive.py from Matplotlib.
+   It's very much subject to change!
 
-- On HTML, will include a .png with a link to a high-res .png.
+Usage
+-----
 
-- On LaTeX, will include a .pdf
+Can be used like this::
 
-This directive supports all of the options of the `image` directive,
-except for `target` (since plot will add its own target).
+    .. plot:: examples/example.py
 
-Additionally, if the :include-source: option is provided, the literal
-source will be included inline, as well as a link to the source.
+    .. plot::
 
-.. warning::
+       import matplotlib.pyplot as plt
+       plt.plot([1,2,3], [4,5,6])
 
-   This is a hacked version of plot_directive.py from Matplotlib.
-   It's very much subject to change!
+    .. plot::
+
+       A plotting example:
+
+       >>> import matplotlib.pyplot as plt
+       >>> plt.plot([1,2,3], [4,5,6])
+
+The content is interpreted as doctest formatted if it has a line starting
+with ``>>>``.
+
+The ``plot`` directive supports the options
+
+    format : {'python', 'doctest'}
+        Specify the format of the input
+    include-source : bool
+        Whether to display the source code. Default can be changed in conf.py
+    
+and the ``image`` directive options ``alt``, ``height``, ``width``,
+``scale``, ``align``, ``class``.
+
+Configuration options
+---------------------
+
+The plot directive has the following configuration options:
+
+    plot_output_dir
+        Directory (relative to config file) where to store plot output.
+        Should be inside the static directory. (Default: 'static')
+
+    plot_pre_code
+        Code that should be executed before each plot.
+
+    plot_rcparams
+        Dictionary of Matplotlib rc-parameter overrides.
+        Has 'sane' defaults.
+
+    plot_include_source
+        Default value for the include-source option
+
+
+TODO
+----
+
+* Don't put temp files to _static directory, but do function in the way
+  the pngmath directive works, and plot figures only during output writing.
+
+* Refactor Latex output; now it's plain images, but it would be nice
+  to make them appear side-by-side, or in floats.
 
 """
 
-import sys, os, glob, shutil, imp, warnings, cStringIO, re
-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 sys, os, glob, shutil, imp, warnings, cStringIO, re, textwrap
+
+def setup(app):
+    setup.app = app
+    setup.config = app.config
+    setup.confdir = app.confdir
+    
+    app.add_config_value('plot_output_dir', '_static', True)
+    app.add_config_value('plot_pre_code', '', True)
+    app.add_config_value('plot_rcparams', sane_rcparameters, True)
+    app.add_config_value('plot_include_source', False, True)
+
+    app.add_directive('plot', plot_directive, True, (1, 0, False),
+                      **plot_directive_options)
+
+sane_rcparameters = {
+    'font.size': 8,
+    'axes.titlesize': 8,
+    'axes.labelsize': 8,
+    'xtick.labelsize': 8,
+    'ytick.labelsize': 8,
+    'legend.fontsize': 8,
+    'figure.figsize': (4, 3),
+}
+
+#------------------------------------------------------------------------------
+# Run code and capture figures
+#------------------------------------------------------------------------------
 
 import matplotlib
 import matplotlib.cbook as cbook
@@ -37,25 +105,41 @@ import matplotlib.pyplot as plt
 import matplotlib.image as image
 from matplotlib import _pylab_helpers
 
-def runfile(fullpath, is_doctest=False):
+def contains_doctest(text):
+    r = re.compile(r'^\s*>>>', re.M)
+    m = r.match(text)
+    return bool(m)
+
+def unescape_doctest(text):
+    """
+    Extract code from a piece of text, which contains either Python code
+    or doctests.
+
+    """
+    if not contains_doctest(text):
+        return text
+
+    code = ""
+    for line in text.split("\n"):
+        m = re.match(r'^\s*(>>>|...) (.*)$', line)
+        if m:
+            code += m.group(2) + "\n"
+        elif line.strip():
+            code += "# " + line.strip() + "\n"
+        else:
+            code += "\n"
+    return code
+
+def run_code(code, code_path):
     # Change the working directory to the directory of the example, so
     # it can get at its data files, if any.
     pwd = os.getcwd()
-    path, fname = os.path.split(fullpath)
-    os.chdir(path)
+    if code_path is not None:
+        os.chdir(os.path.dirname(code_path))
     stdout = sys.stdout
     sys.stdout = cStringIO.StringIO()
     try:
-        code = ""
-        if is_doctest:
-            fd = cStringIO.StringIO()
-            for line in open(fname):
-                m = re.match(r'^\s*(>>>|...) (.*)$', line)
-                if m:
-                    code += m.group(2) + "\n"
-        else:
-            code = open(fname).read()
-
+        code = unescape_doctest(code)
         ns = {}
         exec setup.config.plot_pre_code in ns
         exec code in ns
@@ -64,41 +148,9 @@ def runfile(fullpath, is_doctest=False):
         sys.stdout = stdout
     return ns
 
-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,
-           'doctest-format': directives.flag
-           }
-
-template = """
-.. htmlonly::
-
-   [`source code <%(linkdir)s/%(sourcename)s>`__,
-   `png <%(linkdir)s/%(outname)s.hires.png>`__,
-   `pdf <%(linkdir)s/%(outname)s.pdf>`__]
-
-   .. image:: %(linkdir)s/%(outname)s.png
-%(options)s
-
-.. latexonly::
-   .. image:: %(linkdir)s/%(outname)s.pdf
-%(options)s
-
-"""
-
-exception_template = """
-.. htmlonly::
-
-   [`source code <%(linkdir)s/%(sourcename)s>`__]
-
-Exception occurred rendering plot.
-
-"""
-
+#------------------------------------------------------------------------------
+# Generating figures
+#------------------------------------------------------------------------------
 
 def out_of_date(original, derived):
     """
@@ -108,39 +160,27 @@ def out_of_date(original, derived):
     return (not os.path.exists(derived)
             or os.stat(derived).st_mtime < os.stat(original).st_mtime)
 
-def makefig(fullpath, outdir, is_doctest=False):
+def makefig(code, code_path, output_dir, output_base, config):
     """
     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
-
-    print '    makefig: fullpath=%s, outdir=%s'%( fullpath, outdir)
-    formats = [('png', 80),
+    formats = [('png', 100),
                ('hires.png', 200),
                ('pdf', 50),
                ]
 
-    basedir, fname = os.path.split(fullpath)
-    basename, ext = os.path.splitext(fname)
-    if ext != '.py':
-        basename = fname
-    sourcename = fname
     all_exists = True
 
-    if basedir != outdir:
-        shutil.copyfile(fullpath, os.path.join(outdir, fname))
-
     # Look for single-figure output files first
     for format, dpi in formats:
-        outname = os.path.join(outdir, '%s.%s' % (basename, format))
-        if out_of_date(fullpath, outname):
+        output_path = os.path.join(output_dir, '%s.%s' % (output_base, format))
+        if out_of_date(code_path, output_path):
             all_exists = False
             break
 
     if all_exists:
-        print '    already have %s'%fullpath
         return 1
 
     # Then look for multi-figure output files, assuming
@@ -149,8 +189,9 @@ def makefig(fullpath, outdir, is_doctest=False):
     while True:
         all_exists = True
         for format, dpi in formats:
-            outname = os.path.join(outdir, '%s_%02d.%s' % (basename, i, format))
-            if out_of_date(fullpath, outname):
+            output_path = os.path.join(output_dir,
+                                       '%s_%02d.%s' % (output_base, i, format))
+            if out_of_date(code_path, output_path):
                 all_exists = False
                 break
         if all_exists:
@@ -159,21 +200,23 @@ def makefig(fullpath, outdir, is_doctest=False):
             break
 
     if i != 0:
-        print '    already have %d figures for %s' % (i, fullpath)
         return i
 
     # We didn't find the files, so build them
+    print "-- Plotting figures %s" % output_base
 
-    print '    building %s'%fullpath
-    plt.close('all')    # we need to clear between runs
+    # Clear between runs
+    plt.close('all')
+
+    # Reset figure parameters
     matplotlib.rcdefaults()
-    # Set a figure size that doesn't overflow typical browser windows
-    matplotlib.rcParams['figure.figsize'] = (5.5, 4.5)
+    matplotlib.rcParams.update(config.plot_rcparams)
 
     try:
-        runfile(fullpath, is_doctest=is_doctest)
+        run_code(code, code_path)
     except:
-	s = cbook.exception_to_str("Exception running plot %s" % fullpath)
+        raise
+	s = cbook.exception_to_str("Exception running plot %s" % code_path)
         warnings.warn(s)
         return 0
 
@@ -181,79 +224,193 @@ def makefig(fullpath, outdir, is_doctest=False):
     for i, figman in enumerate(fig_managers):
         for format, dpi in formats:
             if len(fig_managers) == 1:
-                outname = basename
+                name = output_base
             else:
-                outname = "%s_%02d" % (basename, i)
-            outpath = os.path.join(outdir, '%s.%s' % (outname, format))
+                name = "%s_%02d" % (output_base, i)
+            path = os.path.join(output_dir, '%s.%s' % (name, format))
             try:
-                figman.canvas.figure.savefig(outpath, dpi=dpi)
+                figman.canvas.figure.savefig(path, dpi=dpi)
             except:
-                s = cbook.exception_to_str("Exception running plot %s" % fullpath)
+                s = cbook.exception_to_str("Exception running plot %s"
+                                           % code_path)
                 warnings.warn(s)
                 return 0
 
     return len(fig_managers)
 
-def run(arguments, options, state_machine, lineno):
-    reference = directives.uri(arguments[0])
-    basedir, fname = os.path.split(reference)
-    basename, ext = os.path.splitext(fname)
-    if ext != '.py':
-        basename = fname
-    sourcename = fname
-    #print 'plotdir', reference, basename, ext
-
-    # get the directory of the rst file
-    rstdir, rstfile = os.path.split(state_machine.document.attributes['source'])
-    reldir = rstdir[len(setup.confdir)+1:]
-    relparts = [p for p in os.path.split(reldir) if p.strip()]
-    nparts = len(relparts)
-    #print '    rstdir=%s, reldir=%s, relparts=%s, nparts=%d'%(rstdir, reldir, relparts, nparts)
-    #print 'RUN', rstdir, reldir
-    outdir = os.path.join(setup.confdir, setup.config.plot_output_dir, basedir)
-    if not os.path.exists(outdir):
-        cbook.mkdirs(outdir)
-
-    linkdir = ('../' * nparts) + setup.config.plot_output_dir.replace(os.path.sep, '/') + '/' + basedir
-    #linkdir = os.path.join('..', outdir)
-    num_figs = makefig(reference, outdir,
-                       is_doctest=('doctest-format' in options))
-    #print '    reference="%s", basedir="%s", linkdir="%s", outdir="%s"'%(reference, basedir, linkdir, outdir)
-
-    if options.has_key('include-source'):
-        contents = open(reference, 'r').read()
-        if 'doctest-format' in options:
+#------------------------------------------------------------------------------
+# Generating output
+#------------------------------------------------------------------------------
+
+from docutils import nodes, utils
+import jinja
+
+TEMPLATE = """
+{{source_code}}
+
+.. htmlonly::
+
+   {% if source_code %}
+       (`Source code <{{source_link}}>`__)
+   {% endif %}
+
+   .. admonition:: Output
+      :class: plot-output
+
+      {% for name in image_names %}
+      .. figure:: {{link_dir}}/{{name}}.png
+         {%- for option in options %}
+         {{option}}
+         {% endfor %}
+
+         (
+         {%- if not source_code %}`Source code <{{source_link}}>`__, {% endif -%}
+         `PNG <{{link_dir}}/{{name}}.hires.png>`__,
+         `PDF <{{link_dir}}/{{name}}.pdf>`__)
+      {% endfor %}
+
+.. latexonly::
+
+   {% for name in image_names %}
+   .. image:: {{link_dir}}/{{name}}.pdf
+   {% endfor %}
+
+"""
+
+def run(arguments, content, options, state_machine, state, lineno):
+    if arguments and content:
+        raise RuntimeError("plot:: directive can't have both args and content")
+
+    document = state_machine.document
+    config = document.settings.env.config
+
+    options.setdefault('include-source', config.plot_include_source)
+    if options['include-source'] is None:
+        options['include-source'] = config.plot_include_source
+
+    # determine input
+    rst_file = document.attributes['source']
+    rst_dir = os.path.dirname(rst_file)
+    
+    if arguments:
+        file_name = os.path.join(rst_dir, directives.uri(arguments[0]))
+        code = open(file_name, 'r').read()
+        output_base = os.path.basename(file_name)
+    else:
+        file_name = rst_file
+        code = textwrap.dedent("\n".join(map(str, content)))
+        counter = document.attributes.get('_plot_counter', 0) + 1
+        document.attributes['_plot_counter'] = counter
+        output_base = '%d-%s' % (counter, os.path.basename(file_name))
+
+    rel_name = relative_path(file_name, setup.confdir)
+
+    base, ext = os.path.splitext(output_base)
+    if ext in ('.py', '.rst', '.txt'):
+        output_base = base
+
+    # is it in doctest format?
+    is_doctest = contains_doctest(code)
+    if options.has_key('format'):
+        if options['format'] == 'python':
+            is_doctest = False
+        else:
+            is_doctest = True
+
+    # determine output
+    file_rel_dir = os.path.dirname(rel_name)
+    while file_rel_dir.startswith(os.path.sep):
+        file_rel_dir = file_rel_dir[1:]
+
+    output_dir = os.path.join(setup.confdir, setup.config.plot_output_dir,
+                              file_rel_dir)
+
+    if not os.path.exists(output_dir):
+        cbook.mkdirs(output_dir)
+
+    # copy script
+    target_name = os.path.join(output_dir, output_base)
+    f = open(target_name, 'w')
+    f.write(unescape_doctest(code))
+    f.close()
+
+    source_link = relative_path(target_name, rst_dir)
+
+    # determine relative reference
+    link_dir = relative_path(output_dir, rst_dir)
+
+    # make figures
+    num_figs = makefig(code, file_name, output_dir, output_base, config)
+
+    # generate output
+    if options['include-source']:
+        if is_doctest:
             lines = ['']
         else:
             lines = ['.. code-block:: python', '']
-        lines += ['    %s'%row.rstrip() for row in contents.split('\n')]
-        del options['include-source']
+        lines += ['    %s' % row.rstrip() for row in code.split('\n')]
+        source_code = "\n".join(lines)
     else:
-        lines = []
+        source_code = ""
 
-    if 'doctest-format' in options:
-        del options['doctest-format']
-    
     if num_figs > 0:
-        options = ['      :%s: %s' % (key, val) for key, val in
-                   options.items()]
-        options = "\n".join(options)
-
+        image_names = []
         for i in range(num_figs):
             if num_figs == 1:
-                outname = basename
+                image_names.append(output_base)
             else:
-                outname = "%s_%02d" % (basename, i)
-            lines.extend((template % locals()).split('\n'))
+                image_names.append("%s_%02d" % (output_base, i))
     else:
-        lines.extend((exception_template % locals()).split('\n'))
+        reporter = state.memo.reporter
+        sm = reporter.system_message(3, "Exception occurred rendering plot",
+                                     line=lineno)
+        return [sm]
+
+
+    opts = [':%s: %s' % (key, val) for key, val in options.items()
+            if key in ('alt', 'height', 'width', 'scale', 'align', 'class')]
 
+    result = jinja.from_string(TEMPLATE).render(
+        link_dir=link_dir.replace(os.path.sep, '/'),
+        source_link=source_link,
+        options=opts,
+        image_names=image_names,
+        source_code=source_code)
+
+    lines = result.split("\n")
     if len(lines):
         state_machine.insert_input(
             lines, state_machine.input_lines.source(0))
     return []
 
 
+def relative_path(target, base):
+    target = os.path.abspath(os.path.normpath(target))
+    base = os.path.abspath(os.path.normpath(base))
+
+    target_parts = target.split(os.path.sep)
+    base_parts = base.split(os.path.sep)
+    rel_parts = 0
+
+    while target_parts and base_parts and target_parts[0] == base_parts[0]:
+        target_parts.pop(0)
+        base_parts.pop(0)
+
+    rel_parts += len(base_parts)
+    return os.path.sep.join([os.path.pardir] * rel_parts + target_parts)
+
+#------------------------------------------------------------------------------
+# plot:: directive registration etc.
+#------------------------------------------------------------------------------
+
+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
 
 try:
     from docutils.parsers.rst import Directive
@@ -262,34 +419,34 @@ except ImportError:
 
     def plot_directive(name, arguments, options, content, lineno,
                        content_offset, block_text, state, state_machine):
-        return run(arguments, options, state_machine, lineno)
+        return run(arguments, content, options, state_machine, state, 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)
+            return run(self.arguments, self.content, self.options,
+                       self.state_machine, self.state, self.lineno)
     plot_directive.__doc__ = __doc__
 
-    directives.register_directive('plot', plot_directive)
-
-def setup(app):
-    setup.app = app
-    setup.config = app.config
-    setup.confdir = app.confdir
-
-    app.add_config_value('plot_output_dir', '_static', True)
-    app.add_config_value('plot_pre_code', '', True)
-
-plot_directive.__doc__ = __doc__
-
-directives.register_directive('plot', plot_directive)
-
+def _option_boolean(arg):
+    if not arg or not arg.strip():
+        return None
+    elif arg.strip().lower() in ('no', '0', 'false'):
+        return False
+    elif arg.strip().lower() in ('yes', '1', 'true'):
+        return True
+    else:
+        raise ValueError('"%s" unknown boolean' % arg)
+
+def _option_format(arg):
+    return directives.choice(arg, ('python', 'lisp'))
+
+plot_directive_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': _option_boolean,
+                          'format': _option_format,
+                          }