Saving a legend separately with matplotlib

save_legend.py

"""
Here we imagine that you've plotted some things, added a legend, and now
want to be able to save that legend separately.
Perhaps you want to rescale things later, or maybe you've plotted a bunch
of similar graphs and want to have a common legend for all of them, which
can be awkard to do in matplotlib.

Here is one method for doing this:
    0. Plot some things, set up a legend
    1. Get information about the legend, in particular its bounding box
    2. Create a second figure that will only contain the legend
    3. Add the legend to the figure
    4. Save the legend figure

This is not conceptually difficult, the hard thing (for me) was finding
the correct options and commands to produce output that was of the 
proper size, without being cut-off, having a weird DPI, or with an
excessive amount of whitespace.

The code included herein is the process I used, roughly-- this is more of
an example, although it should run, and can be adapted as needed.
"""
import matplotlib as mpl
import matplotlib.pyplot as plt

# Example data
xdata = [i for i in range(10)] 
y1 = [x**2 for x in xdata]
y2 = [x**2 - x for x in xdata]

# Creating the original plot
fig, ax = plt.subplots(figsize=(4,3))

# Plotting with label
line_1 = ax.plot(xdata, y1, label='first line')
line_2 = ax.plot(xdata, y2, label='second line')

# Lines to show on legend and their labels
lines = [line_1, line_2]
labels = [i.get_label() for i in lines]

# Create the legend (that shows up on the main plot)
# These values are changeable, should create a legend outside to the
# right of the graph.
legend = fig.legend(
    lines, 
    labels, 
    loc='upper right',
    bbox_to_anchor=(0.98, 0.9 , 0.32, -0.102),
    mode='expand',
    ncol=1,
    bbox_transform=fig.transFigure,
)

# Save the figure, ensuring that the legend doesn't get cut off
# using `bbox_extra_artists`
figpath = 'graph.png'
fig.savefig(figpath, bbox_extra_artists=[legend], bbox_inches='tight')

# ---------------------------------------------------------------------
# Create a separate figure for the legend
# ---------------------------------------------------------------------
# Bounding box for legend (in order to get proportions right)
# Issuing a draw command seems necessary in order to ensure the layout
# happens correctly; I was getting weird bounding boxes without it.
fig.canvas.draw()
# This gives pixel coordinates for the legend bbox (although perhaps 
# if you are using a different renderer you might get something else).
legend_bbox = legend.get_tightbbox(fig.canvas.get_renderer())
# Convert pixel coordinates to inches
legend_bbox = legend_bbox.transformed(fig.dpi_scale_trans.inverted())

# Create teh separate figure, with appropriate width and height
# Create a separate figure for the legend
legend_fig, legend_ax = plt.subplots(figsize=(legend_bbox.width, legend_bbox.height))

# Recreate the legend on the separate figure/axis
legend_squared = legend_ax.legend(
    *ax.get_legend_handles_labels(), 
    bbox_to_anchor=(0, 0, 1, 1),
    bbox_transform=legend_fig.transFigure,
    frameon=False,
    fancybox=None,
    shadow=False,
    ncol=legend_cols,
    mode='expand',
)

# Remove everything else from the legend's figure
legend_ax.axis('off')

# Save the legend as a separate figure
legend_figpath = 'graph-legend.png'
print(f"Saving to: {legend_figpath}")
legend_fig.savefig(
    legend_figpath,
    bbox_inches='tight',
    bbox_extra_artists=[legend_squared],
)