如何通过自定义类为现有matplotlib Axes实例新增额外方法且保留原有行为

核心实现思路 #

Python是动态类型语言，支持运行时修改实例的所属类，只要自定义类是原Axes类的子类，就可以在完全保留figure.add_subplot创建的实例原有状态的前提下，获得自定义类的所有方法，原生Axes的所有API、类型检查都完全兼容，完美满足你「不手动调用Axes.__init__，基于已有实例扩展功能」的需求。 #

完整可运行代码 #

import matplotlib.axes
import matplotlib.pyplot as plt

class AttentiveAxes(matplotlib.axes.Axes):
    def __new__(cls, figure, geometry):
        # 由figure完成原生Axes实例的创建与全量初始化
        ax = figure.add_subplot(geometry)
        # 将实例的所属类替换为自定义子类，直接获得所有扩展方法
        ax.__class__ = cls
        # 执行自定义初始化逻辑，不会覆盖原有实例的状态
        ax.__init__()
        return ax
    
    def __init__(self):
        self.add_callback(self._on_new_artist_added)
    
    def _on_new_artist_added(self, artist) -> None:
        print(f"Just plotted a new {artist}.")

效果验证 #

运行你给出的预期用法代码：

fig = plt.figure()
ax = AttentiveAxes(fig, 111)
ax.scatter(range(5), range(5))

会正常输出Just plotted a new <matplotlib.collections.PathCollection object at 0x...>，完全符合预期。 #

原错误写法的问题说明 #

1. 第一种__new__直接返回原生Axes实例的写法：你调用cls.__init__(ax)时，传入的self是原生Axes对象，本身没有自定义类中定义的_on_new_artist_added方法，所以会抛出属性不存在的错误。
2. 第二种用__getattr__代理的写法：虽然能勉强实现功能，但存在很多缺陷：isinstance(ax, matplotlib.axes.Axes)检查会返回False，部分依赖类型判断的matplotlib内置逻辑可能异常，也没有IDE类型提示，维护成本很高。

内容的提问来源于stack exchange，提问作者Guimoute