If you just want to save image files and don’t want matplotlib to attempt to create a graphical window you can do this:
import matplotlib
matplotlib.use('Agg') # do this before importing pyplot
from matplotlib import pyplot
pyplot.ioff()
This puts matplotlib into a batch state similar to ROOT’s ROOT.gROOT.SetBatch(True).