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

Question

I am creating two Python scripts to create some graphs for a technical report. In the first script, I define functions that produce graphs from raw data on my hard drive. Each function creates one specific kind of plot that I need. The second script is more like a batch file that should surround these functions and store the resulting graphics on my hard drive.

I need a way to get the plot back in Python. So basically I want to do this:

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

But I don’t know how to plot the variable that I can return. Is it possible? Because?

Answer 1

You can define your charting functions, for example

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 as

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()