How to Return a Plot or Figure in Python Matplotlib?

by

💬 Question: How can I return a Matplotlib axis, plot, or figure object in Python so I can set some defaults and plot a figure quickly based on those defaults?

Basic Solution

To return a Matplotlib object from a function, follow these steps: 👇

  1. Create a function my_plot(x, y, ...) with arguments to customize your plot (e.g., the plot’s x and y data, color, size, style, label, etc.).
  2. Call the plt.subplots() function and assign the resulting figure and axis objects to two variables f and ax.
  3. Plot inside the function using ax.plot(x, y, ...) with the custom arguments using the data from the my_plot() function arguments.
  4. Return the axis object (and possibly the figure object) using return ax (or return f, ax) so that you can further add data to the current plot.

Here’s a concrete example:

import matplotlib.pyplot as plt


def my_plot(x, y, style='o--', label='Data'):
    f, ax = plt.subplots()
    ax.plot(x, y, style, label=label)
    return ax

x = [1, 2, 3]
y1 = [10, 20, 30]
ax = my_plot(x, y1, label='y1')

plt.legend()
plt.grid()
plt.show()

The resulting plot:

In case you struggle with a proper understanding of the plt.subplots() function, I’d recommend you check out the following video tutorial and the associated blog tutorial to keep improving your skills.

🌍 Recommended Tutorial: Matplotlib Subplots – A Helpful Illustrated Guide

The advantage of returning the axis object from the function in the previous code snippet is that you can add more data to the plot later and outside the function like so:

import matplotlib.pyplot as plt


def plot(x, y, style='o--', label='Data'):
    f, ax = plt.subplots()
    ax.plot(x, y, style, label=label)
    return ax


x = [1, 2, 3]
y1 = [10, 20, 30]
ax = plot(x, y1, label='y1')


ax.plot(x, [5, 10, 20], label='y2')
ax.plot(x, [9, 10, 11], label='y3')
ax.plot(x, [15, 15, 15], label='y4')

plt.legend()
plt.grid()
plt.show()

Only the highlighted lines have changed—you can see that this approach of returning an axis object from a Python function gives you maximum flexibility, even outside the function body.

Here’s the output of the previous code snippet with the additional data:

Of course, if you don’t need this, you can also perform all the plotting and housekeeping inside the function and don’t return anything. However, this was not the question addressed in this tutorial, so I’ll skip it for now.

Where to Go From Here?

Make sure to check out our free Matplotlib guide here:

🌍 Recommended Tutorial: Python Matplotlib Ultimate Video Guide

Also, I’d love to see you as an active reader in our Python email community of ambitious learners: