Save Matplotlib subplot as PDF using PyFPDF-CodePudding

I am trying to create two figures with two plots in each figure and then create a PDF file with one of the figures. I want to choose which figure will be saved in the PDF file. I have this simple code. It's just an example:

from fpdf import FPDF
from matplotlib import pyplot as plt
import tempfile
import io

# Examples
X = [0, 1, 2, 3, 4, 5, 6]
Y = [0, 1, 0, 1, 0, 1, 0]

X1 = [0, 1, 2, 3, 4, 6, 7, 8, 9, 10]
Y1 = [1, 4, 6, 4, 2, 12, 4, 6, 3, 12]

X2 = [0, 1, 2, 3, 4, 6, 7, 8, 9, 10]
Y2 = [1, 1, 1, 3, 6, 12, 18, 24, 30, 36]

X3 = [0, 1, 2, 3, 4, 6, 7, 8, 9, 10]
Y3 = [1, 30, 20, 30, 10, 30, 40, 40, 50, 20]


def plot(x,y, x1, y1):

    fig, (ax1, ax2) = plt.subplots(1, 2)
    ax1.plot(x, y)
    ax2.plot(x1, y1)

    return fig

def pdf_generator():

    fig = io.BytesIO()
    plt.savefig(fig, format="png")
    saved_fig = tempfile.NamedTemporaryFile()

    with open(f"{saved_fig.name}.png", 'wb') as sf:
        sf.write(fig.getvalue())

    pdf = FPDF('P', 'mm', 'A4')
    pdf.add_page(orientation='P', format='A4')
    pdf.set_xy(30, 50)
    pdf.image(fig, w=140, h=110)
    fig.close()

    return pdf.output('pdf_plot.pdf')

figA = plot(X1,Y1, X2,Y2)
figB = plot(X,Y, X3,Y3)

pdf_generator()

I created the figA and the figB. However, the figure that is always "converted" to PDF is the figB. I already tried to create a pdf_generator function that would accept as a parameter the figure's name but it didn't work.

Any suggestion to solve this problem?

CodePudding user response：

You're nearly there, you just need to change your function definition to take as an argument the figure you want to plot, and save it using figure.savefig() rather than plt.savefig:

#...
def pdf_generator(figure):
    fig = io.BytesIO()
    figure.savefig(fig, format="png")
    #...

pdf_generator(figA)
pdf_generator(figB)