DOC: adding get_rst.py sphinx extension to build example gallery with mayavi screen capture

Author: agramfort

doc/conf.py

@@ -29,7 +29,7 @@
 # Add any Sphinx extension module names here, as strings. They can be extensions
 # coming with Sphinx (named 'sphinx.ext.*') or your custom ones.
 extensions = ['sphinx.ext.autodoc', 'sphinx.ext.doctest', 'sphinx.ext.coverage',
-              'ipython_console_highlighting']
+              'ipython_console_highlighting', 'gen_rst']
 
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ['_templates']
@@ -47,6 +47,9 @@
 project = u'PySurfer'
 copyright = u'2011, Michael Waskom, Alexandre Gramfort, Scott Burns'
 
+# Generate the plots for the gallery
+plot_gallery = True
+
 # The version info for the project you're documenting, acts as replacement for
 # |version| and |release|, also used in various other places throughout the
 # built documents.

doc/index.rst

@@ -12,11 +12,11 @@ PySurfer offers both a command-line interface designed to broadly
 replicate Freesurfer's Tksurfer program as well as a Python library
 for writing scripts to efficiently explore complex datasets. 
 
-Contents:
-
 .. toctree::
    :maxdepth: 2
 
+   auto_examples/index.rst
+
 Indices and tables
 ==================
 

doc/sphinxext/gen_rst.py

@@ -0,0 +1,205 @@
+"""
+Example generation
+
+Generate the rst files for the examples by iterating over the python
+example files.
+
+Files that generate images should start with 'plot'
+
+"""
+import os
+import shutil
+import traceback
+
+fileList = []
+
+import matplotlib
+matplotlib.use('Agg')
+
+import token, tokenize
+
+rst_template = """
+
+.. _example_%(short_fname)s:
+
+%(docstring)s
+
+**Python source code:** :download:`%(fname)s <%(fname)s>`
+
+.. literalinclude:: %(fname)s
+    :lines: %(end_row)s-
+    """
+
+plot_rst_template = """
+
+.. _example_%(short_fname)s:
+
+%(docstring)s
+
+.. image:: images/%(image_name)s
+    :align: center
+
+**Python source code:** :download:`%(fname)s <%(fname)s>`
+
+.. literalinclude:: %(fname)s
+    :lines: %(end_row)s-
+    """
+
+
+def extract_docstring(filename):
+    """ Extract a module-level docstring, if any
+    """
+    lines = file(filename).readlines()
+    start_row = 0
+    if lines[0].startswith('#!'):
+        lines.pop(0)
+        start_row = 1
+
+    docstring = ''
+    first_par = ''
+    tokens = tokenize.generate_tokens(lines.__iter__().next)
+    for tok_type, tok_content, _, (erow, _), _ in tokens:
+        tok_type = token.tok_name[tok_type]
+        if tok_type in ('NEWLINE', 'COMMENT', 'NL', 'INDENT', 'DEDENT'):
+            continue
+        elif tok_type == 'STRING':
+            docstring = eval(tok_content)
+            # If the docstring is formatted with several paragraphs, extract
+            # the first one:
+            paragraphs = '\n'.join(line.rstrip()
+                                for line in docstring.split('\n')).split('\n\n')
+            if len(paragraphs) > 0:
+                first_par = paragraphs[0]
+        break
+    return docstring, first_par, erow+1+start_row
+
+
+def generate_example_rst(app):
+    """ Generate the list of examples, as well as the contents of
+        examples.
+    """
+    root_dir = os.path.join(app.builder.srcdir, 'auto_examples')
+    example_dir = os.path.abspath(app.builder.srcdir +  '/../' + 'examples')
+    try:
+        plot_gallery = eval(app.builder.config.plot_gallery)
+    except TypeError:
+        plot_gallery = bool(app.builder.config.plot_gallery)
+    if not os.path.exists(example_dir):
+        os.makedirs(example_dir)
+    if not os.path.exists(root_dir):
+        os.makedirs(root_dir)
+
+    # we create an index.rst with all examples
+    fhindex = file(os.path.join(root_dir, 'index.rst'), 'w')
+    fhindex.write("""\
+
+Examples
+========
+
+.. _examples-index:
+""")
+    # Here we don't use an os.walk, but we recurse only twice: flat is
+    # better than nested.
+    generate_dir_rst('.', fhindex, example_dir, root_dir, plot_gallery)
+    for dir in sorted(os.listdir(example_dir)):
+        if dir == '.svn':
+            continue
+        if os.path.isdir(os.path.join(example_dir, dir)):
+            generate_dir_rst(dir, fhindex, example_dir, root_dir, plot_gallery)
+    fhindex.flush()
+
+
+def generate_dir_rst(dir, fhindex, example_dir, root_dir, plot_gallery):
+    """ Generate the rst file for an example directory.
+    """
+    target_dir = os.path.join(root_dir, dir)
+    src_dir = os.path.join(example_dir, dir)
+    if not os.path.exists(os.path.join(src_dir, 'README.txt')):
+        print 80*'_'
+        print ('Example directory %s does not have a README.txt file'
+                        % src_dir)
+        print 'Skipping this directory'
+        print 80*'_'
+        return
+    fhindex.write("""
+
+%s
+
+.. toctree::
+
+""" % file(os.path.join(src_dir, 'README.txt')).read())
+    if not os.path.exists(target_dir):
+        os.makedirs(target_dir)
+    for fname in sorted(os.listdir(src_dir)):
+        if fname.endswith('py'):
+            generate_file_rst(fname, target_dir, src_dir, plot_gallery)
+            fhindex.write('    %s\n' % (os.path.join(dir, fname[:-3])))
+
+
+def generate_file_rst(fname, target_dir, src_dir, plot_gallery):
+    """ Generate the rst file for a given example.
+    """
+    image_name = fname[:-2] + 'png'
+    global rst_template, plot_rst_template
+    this_template = rst_template
+    last_dir = os.path.split(src_dir)[-1]
+    # to avoid leading . in file names
+    if last_dir == '.': last_dir = ''
+    else: last_dir += '_'
+    short_fname =  last_dir + fname
+    src_file = os.path.join(src_dir, fname)
+    example_file = os.path.join(target_dir, fname)
+    shutil.copyfile(src_file, example_file)
+    if plot_gallery and fname.startswith('plot'):
+        # generate the plot as png image if file name
+        # starts with plot and if it is more recent than an
+        # existing image.
+        if not os.path.exists(
+                            os.path.join(target_dir, 'images')):
+            os.makedirs(os.path.join(target_dir, 'images'))
+        image_file = os.path.join(target_dir, 'images', image_name)
+        if (not os.path.exists(image_file) or
+                os.stat(image_file).st_mtime <=
+                    os.stat(src_file).st_mtime):
+            print 'plotting %s' % fname
+            import matplotlib.pyplot as plt
+            plt.close('all')
+            try:
+                from enthought.mayavi import mlab
+                mlab.close(all=True)
+            except:
+                pass
+
+            try:
+                execfile(example_file, {'pl' : plt})
+                facecolor = plt.gcf().get_facecolor() # hack to keep black bg
+                if facecolor == (0.0, 0.0, 0.0, 1.0):
+                    plt.savefig(image_file, facecolor='black')
+                else:
+                    plt.savefig(image_file)
+
+                try:
+                    from enthought.mayavi import mlab
+                    e = mlab.get_engine()
+                    if len(e.scenes) > 0:
+                        mlab.savefig(image_file, size=(400, 400))
+                except:
+                    pass
+
+            except:
+                print 80*'_'
+                print '%s is not compiling:' % fname
+                traceback.print_exc()
+                print 80*'_'
+        this_template = plot_rst_template
+
+    docstring, short_desc, end_row = extract_docstring(example_file)
+
+    f = open(os.path.join(target_dir, fname[:-2] + 'rst'),'w')
+    f.write( this_template % locals())
+    f.flush()
+
+
+def setup(app):
+    app.connect('builder-inited', generate_example_rst)
+    app.add_config_value('plot_gallery', True, 'html')

examples/annotation.py examples/plot_annotation.py

@@ -10,6 +10,8 @@
 from os.path import join as pjoin
 from surfer import viz
 
+###############################################################################
+# Look for files on drive
 subj_dir = os.environ["SUBJECTS_DIR"]
 subject_id = 'fsaverage'
 
@@ -20,7 +22,11 @@
 data_path = pjoin(subj_dir, subject_id)
 annot_path = pjoin(data_path, "label", "%s.aparc.annot" % "lh")
 
+###############################################################################
+# Visualize
 brain = viz.Brain(sub, hemi, surf)
 brain.add_annotation(annot_path, borders=True)
+
+###############################################################################
 # show all views
 brain.show_view('lateral')

examples/plot_basics.py

@@ -0,0 +1,19 @@
+"""
+=======================
+Basics of vizualization
+=======================
+
+"""
+print __doc__
+
+from surfer import Brain
+
+sub = 'fsaverage'
+hemi = 'lh'
+surf = 'inflated'
+
+brain = Brain(sub, hemi, surf)
+
+###############################################################################
+# show all views
+brain.show_view('lateral')