In Python, how can I inherit and override a method on a class instance, assigning this new version to the same name as the old one?

Question:

In Matplotlib, a common problem are unwanted white lines between Patch objects drawn with pcolor, pcolormesh, and contourf (see this question for the former two and this question for the latter).

I’ve attempted to fix this automatically by adding new methods to my Axes class/subclass instances using MethodType. I do this instead of subclassing simply because I want to generate Axes by passing slices of GridSpec objects to the add_subplot method on a Figure instance, and I am not aware of how I could do this with some kind of subclassed matplotlib.axes.Subplot (but I welcome advice). Here is some example code:

from types import MethodType
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

f = plt.figure()
gs = GridSpec(2,1)
ax = f.add_subplot(gs[0,0])

def _pcolormesh(self, *args, **kwargs):
    p = self.pcolormesh(*args, **kwargs)
    p.set_edgecolor('face')
    p.set_linewidth(0.2) # Will cover white lines, without making dot in corner of each square visible
    return p

def _contourf(self, *args, **kwargs):
    c = self.contourf(*args, **kwargs)
    for _ in c.collections:
        _.set_edgecolor('face')
    return c

ax.mypcolormesh = MethodType(_pcolormesh, ax)
ax.mycontourf = MethodType(_contourf, ax)

In the last line, I would prefer to be able to write ax.pcolormesh instead of ax.mypcolormesh, but this raises a RecursionError, because _pcolormesh calls the original method name… which is now aliased to itself.

So, how can I access a method on this Axes instance, override it, and preserve the original name?

Asked By: Luke Davis

||
Source

Answers:

Well… I solved it, sort of. Just use

ax._contourf = ax.contourf

which saves a copy of the old method, and then, for example,

def _contourf(self, *args, **kwargs):
    c = self._contourf(*args, **kwargs)
    for collection in c.collections:
        collection.set_edgecolor('face')
    return c
ax.contourf = MethodType(_contourf, ax)

which assigns the new method to call the copy of the old method.

Answered By: Luke Davis

The efficient solution

Since replacing the method for each axis individually is a lot more typing than using a simple function, the most efficient method would be to create a Python file, myhacks.py, with the respective function

def pcolormesh(ax, *args, **kwargs):
    p = ax.pcolormesh(*args, **kwargs)
    p.set_edgecolor('face')
    p.set_linewidth(0.2)
    return p

And use it whenever the improved version of the pcolormesh is needed:

import matplotlib.pyplot as plt
import myhacks as m
# ...other imports

fig, ax = plt.subplots()
m.pcolormesh(ax, other_arguments)

This works also well for already created files, where one would simply search an replace "ax.pcolormesh(" with "m.pcolormesh(ax," (if necessary using regex for possible other axes names).

The academic solution

It is of course possible to subclass matplotlib.axes.Axes to include the desired function. Since there is no real benefit of this, except knowing how to do it, I would call it "academic solution".

So, again we can create a file, myhacks.py, for our custom class, register the custom class as a projection to Matplotlib,

from matplotlib.axes import Axes
from matplotlib.projections import register_projection

class MyAxes(Axes):
    name = 'mycoolnewaxes'
    def pcolormesh(self,*args, **kwargs):
        p = Axes.pcolormesh(self,*args, **kwargs)
        p.set_edgecolor('face')
        p.set_linewidth(0.2)
        return p

register_projection(MyAxes)

And use this by importing it and creating the axes using the projection:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import myhacks as m

fig = plt.figure()
gs = GridSpec(2,1)
ax = fig.add_subplot(gs[0,0], projection='mycoolnewaxes')

z = np.random.rand(10,13)
ax.pcolormesh(z)

plt.show()
Answered By: ImportanceOfBeingErnest