how to decorate a matplotlib plot with another function

Question

I often have to make plots that build on other existing plots. Instead of writing a new function that copies existing, I wanted to apply a decorator to a matplotlib plot that might add an additional line or such. I wrote the code below which is not working. I was wondering if someone might know the right way to do this.

So here is the goal:

x, y = np.linspace(0,5), np.linspace(0,10)
x2,y2 = np.linspace(0,5), np.linspace(0,5)
plt.plot(x,y)
plt.plot(x2,y2)
plt.show()

Now in the code below I created one plot with a single line. Then I tried to add a second plotted line in the decorator. I know the code does not work, but wanted to see if someone could tell me how to fix this.

def overplot(func):
    def testplot_wrapper(func):
        x, y = np.linspace(0,5), np.linspace(0,10)
        plt.plot(x,y)
        plot.show()
    return(testplot_wrapper)

@overplot
def testplot():
    x, y = np.linspace(0,5), np.linspace(0,10)
    plt.plot(x,y)
    plt.show() 

p = overplot(testplot) 
p.show()

The goal is to produce a plot that

Answer 1

You need to call func:

def overplot(func):
    def testplot_wrapper():
        x, y = np.linspace(0,5), np.linspace(0,10)
        plt.plot(x,y)
        func()
    return(testplot_wrapper)

@overplot
def testplot():
    x, y = np.linspace(0,5), np.linspace(0,5)
    plt.plot(x,y) 

testplot()
plt.show()