Python coroutine for plotting to multiple axes in multiple figures

axgriditer.py

def axgriditer(gridfun=None, figfun=None, **figparams):
    """
    Generator for making a series of plots to multiple axes in multiple figures.

    This function provides a coroutine to handle tasks involved in
    creating a figure with one or more axes and saving the figure when
    all the axes have been used (or all the data has been used),
    repeating as needed.

    Keyword arguments:

    gridfun : callable, or None
      function to open figure and create axes, with signature
         fig,axes = gridfun(**kwargs)
      If None, a default function creates one subplot in a figure.

    figfun : callable or generator

      function called when the figure is full or the generator is
      closed. Should write the figure to disk (if desired) and close
      it.  If a callable, the signature is figfun(fig).  If a
      generator, figfun.send(fig) is called.

    kwargs : passed to gridfun()

    Yields : elements of gridfun()[1] (matplotlib.axes.AxesSubplot objects or equivalent)

    **Examples:**

    import numpy as nx
    axg  = axgriditer( lambda : matplotlib.pyplot.subplots(2,1),
                       lambda fig : fig.show() )
    data = [nx.random.randn(100) for i in range(4)]
    for d,ax in zip(data,axg):
        ax.plot(d)
    axg.close()         # ensure last figure is closed

    Copyright 2009-2012 Daniel Meliza (dan at meliza\org)
    License: Gnu Public License version 3
    """
    if gridfun is None:
        from matplotlib.pyplot import subplots
        gridfun = lambda : subplots(1,1)

    fig,axg = gridfun(**figparams)
    try:
        while 1:
            for ax in axg.flat:
                yield ax
            if callable(figfun): figfun(fig)
            elif hasattr(figfun,'send'): figfun.send(fig)
            fig, axg = gridfun(**figparams)
    except:
        ## catches StopIteration as well as other "real" exceptions.
        ## only calls figfun if there are unsaved axes
        if fig and hasattr(fig,'axes') and len(fig.axes):
            if callable(figfun): figfun(fig)
            elif hasattr(figfun,'send'): figfun.send(fig)
        raise