Pylab Save/Plot Decorator

save_show.py

import matplotlib.pyplot as plt

def save_show_no(plotfunc):
    # A function decorator that adds the option to save or show a plot
    # depending on whether a filename option is set.

    def decorate(*args,**kwargs):

        ax = plotfunc(*args)

        if 'filename' in kwargs.keys():

            plt.savefig(kwargs['filename'])

        elif 'show' in kwargs.keys():

            plt.show()

        else:

            return ax

    return decorate