sphinxext: revise plot_directive

1 files changed, 325 insertions, 250 deletions

diff --git a/doc/sphinxext/plot_directive.py b/doc/sphinxext/plot_directive.py
index 45b4c3f78..b33da37e2 100644
--- a/doc/sphinxext/plot_directive.py
+++ b/doc/sphinxext/plot_directive.py
@@ -6,6 +6,7 @@ A special directive for generating a matplotlib plot.
    This is a hacked version of plot_directive.py from Matplotlib.
    It's very much subject to change!
 
+
 Usage
 -----
 
@@ -32,6 +33,7 @@ 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
     
@@ -43,36 +45,34 @@ 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_include_source
+        Default value for the include-source option
 
     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
+    plot_basedir
+        Base directory, to which plot:: file names are relative to.
+        (If None or empty, file names are relative to the directoly where
+        the file containing the directive is.)
 
     plot_formats
-        The set of files to generate. Default: ['png', 'pdf', 'hires.png'],
-        ie. everything.
+        File formats to generate. List of tuples or strings::
+
+            [(suffix, dpi), suffix, ...]
+
+        that determine the file format and the DPI. For entries whose
+        DPI was omitted, sensible defaults are chosen.
 
 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, textwrap
+import sys, os, glob, shutil, imp, warnings, cStringIO, re, textwrap, traceback
 
 import warnings
 warnings.warn("A plot_directive module is also available under "
@@ -81,48 +81,281 @@ warnings.warn("A plot_directive module is also available under "
               "integrated there.",
               FutureWarning, stacklevel=2)
 
+
+#------------------------------------------------------------------------------
+# Registration hook
+#------------------------------------------------------------------------------
+
 def setup(app):
     setup.app = app
     setup.config = app.config
     setup.confdir = app.confdir
-
-    static_path = '_static'
-    if hasattr(app.config, 'html_static_path') and app.config.html_static_path:
-        static_path = app.config.html_static_path[0]
-
-    app.add_config_value('plot_output_dir', static_path, 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_config_value('plot_formats', ['png', 'hires.png', 'pdf'], True)
+    app.add_config_value('plot_basedir', None, True)
 
     app.add_directive('plot', plot_directive, True, (0, 1, False),
                       **plot_directive_options)
 
-sane_rcparameters = {
-    'font.size': 9,
-    'axes.titlesize': 9,
-    'axes.labelsize': 9,
-    'xtick.labelsize': 9,
-    'ytick.labelsize': 9,
-    'legend.fontsize': 9,
-    'figure.figsize': (4, 3),
-}
+#------------------------------------------------------------------------------
+# plot:: directive
+#------------------------------------------------------------------------------
+from docutils.parsers.rst import directives
+from docutils import nodes
+
+def plot_directive(name, arguments, options, content, lineno,
+                   content_offset, block_text, state, state_machine):
+    return run(arguments, content, options, state_machine, state, lineno)
+plot_directive.__doc__ = __doc__
+
+def _option_boolean(arg):
+    if not arg or not arg.strip():
+        # no argument given, assume used as a flag
+        return True
+    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'))
+
+def _option_align(arg):
+    return directives.choice(arg, ("top", "middle", "bottom", "left", "center",
+                                   "right"))
+
+plot_directive_options = {'alt': directives.unchanged,
+                          'height': directives.length_or_unitless,
+                          'width': directives.length_or_percentage_or_unitless,
+                          'scale': directives.nonnegative_int,
+                          'align': _option_align,
+                          'class': directives.class_option,
+                          'include-source': _option_boolean,
+                          'format': _option_format,
+                          }
+
+#------------------------------------------------------------------------------
+# Generating output
+#------------------------------------------------------------------------------
+
+from docutils import nodes, utils
+
+try:
+    # Sphinx depends on either Jinja or Jinja2
+    import jinja2
+    def format_template(template, **kw):
+        return jinja2.Template(template).render(**kw)
+except ImportError:
+    import jinja
+    def format_template(template, **kw):
+        return jinja.from_string(template, **kw)
+
+TEMPLATE = """
+{{ source_code }}
+
+.. htmlonly::
+
+   {% if source_code %}
+       (`Source code <{{ source_link }}>`__)
+
+   .. admonition:: Output
+      :class: plot-output
+
+   {% endif %}
+
+       {% for img in images %}
+       .. figure:: {{ build_dir }}/{{ img.basename }}.png
+          {%- for option in options %}
+          {{ option }}
+          {% endfor %}
+    
+          (
+          {%- if not source_code -%}
+            `Source code <{{source_link}}>`__
+            {%- for fmt in img.formats -%} 
+            , `{{ fmt }} <{{ dest_dir }}/{{ img.basename }}.{{ fmt }}>`__
+            {%- endfor -%}
+          {%- else -%}
+            {%- for fmt in img.formats -%} 
+            {%- if not loop.first -%}, {% endif -%}
+            `{{ fmt }} <{{ dest_dir }}/{{ img.basename }}.{{ fmt }}>`__
+            {%- endfor -%}
+          {%- endif -%}
+          )
+       {% endfor %}
+
+.. latexonly::
+
+   {% for img in images %}
+   .. image:: {{ build_dir }}/{{ img.basename }}.pdf
+   {% endfor %}
+
+"""
+
+class ImageFile(object):
+    def __init__(self, basename, dirname):
+        self.basename = basename
+        self.dirname = dirname
+        self.formats = []
+
+    def filename(self, format):
+        return os.path.join(self.dirname, "%s.%s" % (self.basename, format))
+
+    def filenames(self):
+        return [self.filename(fmt) for fmt in self.formats]
+
+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)
+
+    # determine input
+    rst_file = document.attributes['source']
+    rst_dir = os.path.dirname(rst_file)
+
+    if arguments:
+        if not config.plot_basedir:
+            source_file_name = os.path.join(rst_dir,
+                                            directives.uri(arguments[0]))
+        else:
+            source_file_name = os.path.join(setup.confdir, config.plot_basedir,
+                                            directives.uri(arguments[0]))
+        code = open(source_file_name, 'r').read()
+        output_base = os.path.basename(source_file_name)
+    else:
+        source_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
+        base, ext = os.path.splitext(os.path.basename(source_file_name))
+        output_base = '%s-%d.py' % (base, counter)
+
+    base, source_ext = os.path.splitext(output_base)
+    if source_ext in ('.py', '.rst', '.txt'):
+        output_base = base
+    else:
+        source_ext = ''
+
+    # ensure that LaTeX includegraphics doesn't choke in foo.bar.pdf filenames
+    output_base = output_base.replace('.', '-')
+
+    # 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 directory name fragment
+    source_rel_name = relpath(source_file_name, setup.confdir)
+    source_rel_dir = os.path.dirname(source_rel_name)
+    while source_rel_dir.startswith(os.path.sep):
+        source_rel_dir = source_rel_dir[1:]
+
+    # build_dir: where to place output files (temporarily)
+    build_dir = os.path.join(os.path.dirname(setup.app.doctreedir),
+                             'plot_directive',
+                             source_rel_dir)
+    if not os.path.exists(build_dir):
+        os.makedirs(build_dir)
+
+    # output_dir: final location in the builder's directory
+    dest_dir = os.path.abspath(os.path.join(setup.app.builder.outdir,
+                                            source_rel_dir))
+
+    # how to link to files from the RST file
+    dest_dir_link = os.path.join(relpath(setup.confdir, rst_dir),
+                                 source_rel_dir).replace(os.path.sep, '/')
+    build_dir_link = relpath(build_dir, rst_dir).replace(os.path.sep, '/')
+    source_link = dest_dir_link + '/' + output_base + source_ext
+
+    # make figures
+    try:
+        images = makefig(code, source_file_name, build_dir, output_base,
+                         config)
+    except PlotError, err:
+        reporter = state.memo.reporter
+        sm = reporter.system_message(
+            3, "Exception occurred in plotting %s: %s" % (output_base, err),
+            line=lineno)
+        return [sm]
+
+    # generate output restructuredtext
+    if options['include-source']:
+        if is_doctest:
+            lines = ['']
+        else:
+            lines = ['.. code-block:: python', '']
+        lines += ['    %s' % row.rstrip() for row in code.split('\n')]
+        source_code = "\n".join(lines)
+    else:
+        source_code = ""
+
+    opts = [':%s: %s' % (key, val) for key, val in options.items()
+            if key in ('alt', 'height', 'width', 'scale', 'align', 'class')]
+
+    result = format_template(
+        TEMPLATE,
+        dest_dir=dest_dir_link,
+        build_dir=build_dir_link,
+        source_link=source_link,
+        options=opts,
+        images=images,
+        source_code=source_code)
+
+    lines = result.split("\n")
+    if len(lines):
+        state_machine.insert_input(
+            lines, state_machine.input_lines.source(0))
+
+    # copy image files to builder's output directory
+    if not os.path.exists(dest_dir):
+        os.makedirs(dest_dir)
+
+    for img in images:
+        for fn in img.filenames():
+            shutil.copyfile(fn, os.path.join(dest_dir, os.path.basename(fn)))
+
+    # copy script (if necessary)
+    if source_file_name == rst_file:
+        target_name = os.path.join(dest_dir, output_base + source_ext)
+        f = open(target_name, 'w')
+        f.write(unescape_doctest(code))
+        f.close()
+
+    return []
+
 
 #------------------------------------------------------------------------------
 # Run code and capture figures
 #------------------------------------------------------------------------------
 
 import matplotlib
-import matplotlib.cbook as cbook
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 import matplotlib.image as image
 from matplotlib import _pylab_helpers
 
+import exceptions
+
 def contains_doctest(text):
+    try:
+        # check if it's valid Python as-is
+        compile(text, '<string>', 'exec')
+        return False
+    except SyntaxError:
+        pass
     r = re.compile(r'^\s*>>>', re.M)
-    m = r.match(text)
+    m = r.search(text)
     return bool(m)
 
 def unescape_doctest(text):
@@ -136,7 +369,7 @@ def unescape_doctest(text):
 
     code = ""
     for line in text.split("\n"):
-        m = re.match(r'^\s*(>>>|...) (.*)$', line)
+        m = re.match(r'^\s*(>>>|\.\.\.) (.*)$', line)
         if m:
             code += m.group(2) + "\n"
         elif line.strip():
@@ -145,6 +378,9 @@ def unescape_doctest(text):
             code += "\n"
     return code
 
+class PlotError(RuntimeError):
+    pass
+
 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.
@@ -158,14 +394,22 @@ def run_code(code, code_path):
     # Redirect stdout
     stdout = sys.stdout
     sys.stdout = cStringIO.StringIO()
+
+    # Reset sys.argv
+    old_sys_argv = sys.argv
+    sys.argv = [code_path]
     
     try:
-        code = unescape_doctest(code)
-        ns = {}
-        exec setup.config.plot_pre_code in ns
-        exec code in ns
+        try:
+            code = unescape_doctest(code)
+            ns = {}
+            exec setup.config.plot_pre_code in ns
+            exec code in ns
+        except exceptions.BaseException, err:
+            raise PlotError(traceback.format_exc())
     finally:
         os.chdir(pwd)
+        sys.argv = old_sys_argv
         sys.path[:] = old_sys_path
         sys.stdout = stdout
     return ns
@@ -186,218 +430,91 @@ def out_of_date(original, derived):
 
 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
+    Run a pyplot script *code* and save the images under *output_dir*
+    with file names derived from *output_base*
 
     """
 
-    included_formats = config.plot_formats
-    if type(included_formats) is str:
-        included_formats = eval(included_formats)
-    
-    formats = [x for x in [('png', 80), ('hires.png', 200), ('pdf', 50)]
-               if x[0] in config.plot_formats]
+    # -- Parse format list
+    default_dpi = {'png': 80, 'hires.png': 200, 'pdf': 50}
+    formats = []
+    for fmt in config.plot_formats:
+        if isinstance(fmt, str):
+            formats.append((fmt, default_dpi.get(fmt, 80)))
+        elif type(fmt) in (tuple, list) and len(fmt)==2:
+            formats.append((str(fmt[0]), int(fmt[1])))
+        else:
+            raise PlotError('invalid image format "%r" in plot_formats' % fmt)
 
-    all_exists = True
+    # -- Try to determine if all images already exist
 
     # Look for single-figure output files first
+    all_exists = True
+    img = ImageFile(output_base, output_dir)
     for format, dpi in formats:
-        output_path = os.path.join(output_dir, '%s.%s' % (output_base, format))
-        if out_of_date(code_path, output_path):
+        if out_of_date(code_path, img.filename(format)):
             all_exists = False
             break
+        img.formats.append(format)
 
     if all_exists:
-        return [output_base]
+        return [img]
 
     # Then look for multi-figure output files
-    image_names = []
+    images = []
+    all_exists = True
     for i in xrange(1000):
-        image_names.append('%s_%02d' % (output_base, i))
+        img = ImageFile('%s_%02d' % (output_base, i), output_dir)
         for format, dpi in formats:
-            output_path = os.path.join(output_dir,
-                                       '%s.%s' % (image_names[-1], format))
-            if out_of_date(code_path, output_path):
+            if out_of_date(code_path, img.filename(format)):
                 all_exists = False
                 break
+            img.formats.append(format)
+        
+        # assume that if we have one, we have them all
         if not all_exists:
-            # assume that if we have one, we have them all
             all_exists = (i > 0)
             break
+        images.append(img)
 
     if all_exists:
-        return image_names
+        return images
 
-    # We didn't find the files, so build them
-    print "-- Plotting figures %s" % output_base
+    # -- We didn't find the files, so build them
 
     # Clear between runs
     plt.close('all')
 
-    # Reset figure parameters
-    matplotlib.rcdefaults()
-    matplotlib.rcParams.update(config.plot_rcparams)
-
     # Run code
     run_code(code, code_path)
 
     # Collect images
-    image_names = []
+    images = []
 
     fig_managers = _pylab_helpers.Gcf.get_all_fig_managers()
     for i, figman in enumerate(fig_managers):
         if len(fig_managers) == 1:
-            name = output_base
+            img = ImageFile(output_base, output_dir)
         else:
-            name = "%s_%02d" % (output_base, i)
-        image_names.append(name)
+            img = ImageFile("%s_%02d" % (output_base, i), output_dir)
+        images.append(img)
         for format, dpi in formats:
-            path = os.path.join(output_dir, '%s.%s' % (name, format))
-            figman.canvas.figure.savefig(path, dpi=dpi)
+            try:
+                figman.canvas.figure.savefig(img.filename(format), dpi=dpi)
+            except exceptions.BaseException, err:
+                raise PlotError(traceback.format_exc())
+            img.formats.append(format)
+
+    return images
 
-    return image_names
 
 #------------------------------------------------------------------------------
-# Generating output
+# Relative pathnames
 #------------------------------------------------------------------------------
 
-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 = relpath(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 = relpath(target_name, rst_dir)
-
-    # determine relative reference
-    link_dir = relpath(output_dir, rst_dir)
-
-    # make figures
-    try:
-        image_names = makefig(code, file_name, output_dir, output_base, config)
-    except RuntimeError, err:
-        reporter = state.memo.reporter
-        sm = reporter.system_message(3, "Exception occurred rendering plot",
-                                     line=lineno)
-        return [sm]
-
-    # generate output
-    if options['include-source']:
-        if is_doctest:
-            lines = ['']
-        else:
-            lines = ['.. code-block:: python', '']
-        lines += ['    %s' % row.rstrip() for row in code.split('\n')]
-        source_code = "\n".join(lines)
-    else:
-        source_code = ""
-
-    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 []
-
-
-if hasattr(os.path, 'relpath'):
-    relpath = os.path.relpath
-else:
+try:
+    from os.path import relpath
+except ImportError:
     def relpath(target, base=os.curdir):
         """
         Return a relative path to the target from either the current
@@ -433,45 +550,3 @@ else:
 
         rel_list = [os.pardir] * (len(base_list)-i) + target_list[i:]
         return os.path.join(*rel_list)
-
-#------------------------------------------------------------------------------
-# 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
-
-def plot_directive(name, arguments, options, content, lineno,
-                   content_offset, block_text, state, state_machine):
-    return run(arguments, content, options, state_machine, state, lineno)
-
-plot_directive.__doc__ = __doc__
-
-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,
-                          }