How to assign a plot to a variable and use the variable as the return value in a Python function

Question

I am creating two Python scripts to produce some plots for a technical report. In the first script I am defining functions that produce plots from raw data on my hard-disk. Each function produces one specific kind of plot that I need. The second script is more like a batch file which is supposed to loop around those functions and store the produced plots on my hard-disk.

What I need is a way to return a plot in Python. So basically I want to do this:

fig = some_function_that_returns_a_plot(args)
fig.savefig('plot_name')

But what I do not know is how to make a plot a variable that I can return. Is this possible? Is so, how?

Answer 1

You can define your plotting functions like

import numpy as np
import matplotlib.pyplot as plt

# an example graph type
def fig_barh(ylabels, xvalues, title=''):
    # create a new figure
    fig = plt.figure()

    # plot to it
    yvalues = 0.1 + np.arange(len(ylabels))
    plt.barh(yvalues, xvalues, figure=fig)
    yvalues += 0.4
    plt.yticks(yvalues, ylabels, figure=fig)
    if title:
        plt.title(title, figure=fig)

    # return it
    return fig

then use them like

from matplotlib.backends.backend_pdf import PdfPages

def write_pdf(fname, figures):
    doc = PdfPages(fname)
    for fig in figures:
        fig.savefig(doc, format='pdf')
    doc.close()

def main():
    a = fig_barh(['a','b','c'], [1, 2, 3], 'Test #1')
    b = fig_barh(['x','y','z'], [5, 3, 1], 'Test #2')
    write_pdf('test.pdf', [a, b])

if __name__=="__main__":
    main()