How to plot square function with matplotlib

I have an list of values that alternate between 0 and 1, eg [0,1,0,1,0] and I want to graph them so they appear as a square wave using matplotlib for python. I have this so far:

input_amp = [1,0,1,0,1,0,1,0,1,0]
plt.plot(input_amp, marker='d', color='blue')
plt.title("Waveform")
plt.ylabel('Amplitude')
plt.xlabel("Time")
plt.savefig("waveform.png")
plt.show()

This gives me an output like this this :

How do I make it so instead of going on an angle between the points the line stays flat?

I found this post but it deals more with an animation and not just plotting the function.

Solution

The relevant bit from that post you reference is the drawstyle:

 plt.plot(input_amp, marker='d', color='blue', drawstyle='steps-pre')