""" 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, 'examples') example_dir = os.path.abspath(app.builder.srcdir + '/../' + 'examples') doc_data_link = os.path.join(app.builder.srcdir, 'example_data') example_data_dir = os.path.join(example_dir, 'example_data') 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) if not os.path.exists(doc_data_link): os.symlink(example_data_dir, doc_data_link) # 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, 'index.rst')): print 80*'_' print ('Example directory %s does not have a index.rst file' % src_dir) print 'Skipping this directory' print 80*'_' return fhindex.write(""" %s .. toctree:: """ % file(os.path.join(src_dir, 'index.rst')).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: try: from mayavi import mlab except ImportError: 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: try: from mayavi import mlab except ImportError: 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')