子类化np.ndarray时如何保留原方法返回类型而非子类实例？

我按照numpy官方子类化指南创建了继承自np.ndarray的InfoArray类：

class InfoArray(np.ndarray):
    def __new__(subtype, shape, dtype=float, buffer=None, offset=0,
                strides=None, order=None, info=None):
        obj = super().__new__(subtype, shape, dtype,
                              buffer, offset, strides, order)
        obj.info = info
        return obj
    def __array_finalize__(self, obj):
        self.info = getattr(obj, 'info', None)

但调用该类实例的ndarray方法（如mean）时，返回的是InfoArray实例而非原生类型：

a = InfoArray(shape=(5,), info="some information")
a.mean()
>>> InfoArray(0.5)
type(a.mean())
>>> __main__.InfoArray

而原生np.array的mean方法返回np.float64类型：

b = np.array([1, 2, 3])
b.mean()
>>> 2.0
type(b.mean())
>>> np.float64

我可以通过重写单个方法来解决，但希望避免逐个重写所有方法，请问该如何处理？

核心是重写__array_wrap__方法——这是numpy用来控制操作后返回对象类型的钩子。通过在该方法中判断结果是否为0维数组（标量），我们可以统一处理所有方法的返回类型：

class InfoArray(np.ndarray):
    def __new__(subtype, shape, dtype=float, buffer=None, offset=0,
                strides=None, order=None, info=None):
        obj = super().__new__(subtype, shape, dtype,
                              buffer, offset, strides, order)
        obj.info = info
        return obj
    
    def __array_finalize__(self, obj):
        self.info = getattr(obj, 'info', None)
    
    def __array_wrap__(self, arr, context=None):
        # 若结果是0维数组（标量），返回原生numpy标量类型
        if arr.shape == ():
            return arr.item()
        # 否则返回InfoArray实例，并保留原实例的info属性
        result = super().__array_wrap__(self, arr, context)
        result.info = self.info
        return result

测试验证 #

a = InfoArray(shape=(5,), info="some information")
# 标量返回场景
mean_val = a.mean()
print(mean_val)  # 0.5
print(type(mean_val))  # <class 'numpy.float64'>

# 数组返回场景
sum_arr = a.sum(axis=0, keepdims=True)
print(type(sum_arr))  # <class '__main__.InfoArray'>
print(sum_arr.info)  # some information

原理说明 #

• __array_wrap__会在numpy内置方法（如mean、sum）返回数组对象前被调用，负责包装最终结果。
• 当方法返回标量时，numpy内部会生成一个0维ndarray，通过shape == ()识别后，调用item()转换为原生numpy标量。
• 对于返回数组的操作，继续返回InfoArray实例并继承原实例的info属性，保证子类特性不丢失。

如果需要更精细的控制（比如某些数组操作也想返回原生ndarray），可以结合context参数（包含操作的元信息，如调用的函数名）做针对性判断。

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