Let’s modify the tax equilibrium formula we used before, we just need to add 2 new triangles to represent consumer and producer surplus. The idea is the same, we need to draw the horizontal line, except it is now two horizontal lines, one on consumer price and one on producer price

triangle3 = patches.Polygon([[nonTaxQ,nonTaxPrice],[equilibriumQ,nonTaxPrice],[equilibriumQ,equilibriumP1]],True,color="orange")
triangle4 = patches.Polygon([[nonTaxQ,nonTaxPrice],[equilibriumQ,nonTaxPrice],[equilibriumQ,equilibriumP2]],True,color="purple")
currentAxis.add_patch(triangle3)
currentAxis.add_patch(triangle4)

Also we will add the following print outs based on the area of each piece, and update the legend.

print("The equilibrium price without taxes would be "+str(nonTaxPrice)+" and the quantity would be "+str(nonTaxQ))
print("The consumer surplus without taxes is equal to "+str((priceEnd-nonTaxPrice)*.5*nonTaxQ))
print("The producer surplus without taxes is equal to "+str((nonTaxPrice)*.5*nonTaxQ))
print("")
print("The equilibrium prices are "+str(equilibriumP1)+" and "+str(equilibriumP2)+" and equilibrium quantity is "+str(equilibriumQ)+".")
print("Taxes raised from consumers equals "+str(equilibriumQ*(equilibriumP1-nonTaxPrice)))
print("Taxes raised from producers equals "+str(equilibriumQ*(nonTaxPrice-equilibriumP2)))
print("The deadweight loss to consumers is equal to "+str((equilibriumP1-nonTaxPrice)*.5*(nonTaxQ-equilibriumQ)))
print("The deadweight loss to producers is equal to "+str((nonTaxPrice-equilibriumP2)*.5*(nonTaxQ-equilibriumQ)))
print("")
print("The new consumer surplus is "+str(equilibriumQ*(priceEnd-equilibriumP1)*.5))
print("The new producer surplus is "+str(equilibriumQ*(equilibriumP2)*.5))

We will define the equation solving function again, and then we have our final function:

import sympy
p = sympy.Symbol("p")
q = sympy.Symbol("q")
cTax = sympy.Symbol("cTax")
pTax = sympy.Symbol("pTax")
def eqSolve(eq1,eq2,tax):
    demandP = sympy.solve(eq1-q,p)[0]
    supplyP = sympy.solve(eq2-q,p)[0]
    demandP = demandP-cTax
    supplyP = supplyP+pTax

    demandQ = sympy.solve(demandP-p,q)[0]
    supplyQ = sympy.solve(supplyP-p,q)[0]

    return sympy.solve((demandP-supplyP, demandQ-supplyQ,tax-cTax-pTax), q,p,cTax,pTax)[q]
def EquilibriumTax(demandEquation,supplyEquation,priceStart,tax):
    priceEnd = sympy.solve(demandEquation)[0]
    prices = []
    demand = []
    supply = []
    for price in range(priceStart,priceEnd+1):
        prices += [price]
        demand += [demandEquation.subs(p,price)]
        supply += [supplyEquation.subs(p,price)]



    nonTaxPrice = sympy.solve(demandEquation-supplyEquation)[0]
    nonTaxQ = demandEquation.subs(p,nonTaxPrice)


    equilibriumQ = eqSolve(demandEquation,supplyEquation,tax)
    equilibriumP1 = sympy.solve(demandEquation-equilibriumQ)[0]
    equilibriumP2 = sympy.solve(supplyEquation-equilibriumQ)[0]

    triangle1 = patches.Polygon([[nonTaxQ,nonTaxPrice],[equilibriumQ,nonTaxPrice],[equilibriumQ,equilibriumP1]],True,color="green")
    triangle2 = patches.Polygon([[nonTaxQ,nonTaxPrice],[equilibriumQ,nonTaxPrice],[equilibriumQ,equilibriumP2]],True)
    currentAxis = plt.gca()
    currentAxis.add_patch(triangle1)
    currentAxis.add_patch(triangle2)


    rect1 = patches.Rectangle((0,nonTaxPrice),equilibriumQ,equilibriumP1-nonTaxPrice,linewidth=1,facecolor="red")
    rect2 = patches.Rectangle((0,nonTaxPrice),equilibriumQ,equilibriumP2-nonTaxPrice,linewidth=1,facecolor="yellow")
    currentAxis.add_patch(rect1)
    currentAxis.add_patch(rect2)

    triangle3 = patches.Polygon([[nonTaxQ,nonTaxPrice],[equilibriumQ,nonTaxPrice],[equilibriumQ,equilibriumP1]],True,color="orange")
    triangle4 = patches.Polygon([[nonTaxQ,nonTaxPrice],[equilibriumQ,nonTaxPrice],[equilibriumQ,equilibriumP2]],True,color="purple")
    currentAxis.add_patch(triangle3)
    currentAxis.add_patch(triangle4)


    plt.plot(demand,prices)
    plt.plot(supply,prices)


    plt.legend([rect1,rect2,triangle1,triangle2,triangle3,triangle4], ["Consumer Tax","Producer Tax","Consumer Deadweight Loss","Producer Deadweight Loss","Consumer Surplus","Producer Surplus"])
    plt.plot(equilibriumQ,equilibriumP1, 'ro')
    plt.plot(equilibriumQ,equilibriumP2, 'ro')
    plt.xlabel("Supply and Demand Quantity")
    plt.ylabel("Price")
    plt.show()
    print("The equilibrium price without taxes would be "+str(nonTaxPrice)+" and the quantity would be "+str(nonTaxQ))
    print("The consumer surplus without taxes is equal to "+str((priceEnd-nonTaxPrice)*.5*nonTaxQ))
    print("The producer surplus without taxes is equal to "+str((nonTaxPrice)*.5*nonTaxQ))
    print("")
    print("The equilibrium prices are "+str(equilibriumP1)+" and "+str(equilibriumP2)+" and equilibrium quantity is "+str(equilibriumQ)+".")
    print("Taxes raised from consumers equals "+str(equilibriumQ*(equilibriumP1-nonTaxPrice)))
    print("Taxes raised from producers equals "+str(equilibriumQ*(nonTaxPrice-equilibriumP2)))
    print("The deadweight loss to consumers is equal to "+str((equilibriumP1-nonTaxPrice)*.5*(nonTaxQ-equilibriumQ)))
    print("The deadweight loss to producers is equal to "+str((nonTaxPrice-equilibriumP2)*.5*(nonTaxQ-equilibriumQ)))
    print("")
    print("The new consumer surplus is "+str(equilibriumQ*(priceEnd-equilibriumP1)*.5))
    print("The new producer surplus is "+str(equilibriumQ*(equilibriumP2)*.5))

Now the fun part, plot some graphs!

EquilibriumTax(10-p,p,0,0)
EquilibriumTax(10-p,p,0,2)
EquilibriumTax(15-p,p*2,0,4)

Source Code