Extending matplotlib Line2D class

Question

In my answer to this question I tried to extend Line2D.draw method so that it can be used with plot.

So I changed matplotlib.lines.Line2D to my Line2D extension:

import matplotlib.lines
import matplotlib.pyplot as plt
import matplotlib.patches as pt


class Line2D(matplotlib.lines.Line2D):
    def draw(self, rdr):
        super().draw(rdr)
        xy = self.get_xydata()
        start, end = xy[0], xy[-1]
        r = pt.Rectangle((start[0] - .05, start[1] - .05), .1, .1,
                         color=self.get_color(),
                         fill=None)
        plt.gca().add_patch(r)
        c = pt.Ellipse(end, .05, .05,
                       color=self.get_color(),
                       fill=True)
        plt.gca().add_patch(c)


fig = plt.figure()
ax = fig.add_subplot()
rg = [-.5, 1.5]
ax.set_xlim(rg)
ax.set_ylim(rg)


original_line2D = matplotlib.lines.Line2D
matplotlib.lines.Line2D = Line2D

plt.plot([0, 1], [1, 0], "r--")

plt.show()

I got this result, and I guess plt.plot is used internally... but how to avoid these ghosts drawings?

Answer 1

I think you are allowed to use your line, but then you should reset it to matplotlib's original Line2D. For example:

fig = plt.figure()
ax = fig.add_subplot()
rg = [-.5, 1.5]
ax.set_xlim(rg)
ax.set_ylim(rg)

original_line2D = matplotlib.lines.Line2D
matplotlib.lines.Line2D = MyLine2D

plt.plot([0, 1], [1, 0], "r--")
matplotlib.lines.Line2D = original_line2D
plt.plot([1, 0], [1, 1], "g")
plt.show()

Removing the last line (green), works as expected.