在NumPy ndarray子类的自定义ufunc中如何主动抛出异常？

问题根源 #

当__array_ufunc__中调用np.equal抛出异常后，Python会自动回退到默认的对象相等性比较（__eq__），而NumPy旧版本处理形状不匹配的数组比较时，旧行为是返回False并发出DeprecationWarning，而非直接抛出异常。

解决方案：重写__eq__方法 #

直接重写__eq__方法，绕过默认回退逻辑，强制在形状不匹配时抛出异常，确保自定义equal逻辑的异常被正确传播。

修改后的完整代码：

import numpy as np

class MyArray(np.ndarray):

    def __new__(cls, data):
        return np.array(data).view(MyArray)

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        # 仅转换MyArray类型的输入为原生ndarray
        inputs = [i.view(np.ndarray) if isinstance(i, MyArray) else i for i in inputs]
        if ufunc == np.equal and method == "__call__":
            # 保持返回结果为MyArray类型
            return self._custom_equal(*inputs, **kwargs).view(MyArray)
        return super().__array_ufunc__(ufunc, method, *inputs, **kwargs)

    def __eq__(self, other):
        # 统一转换输入为原生ndarray
        self_arr = self.view(np.ndarray)
        other_arr = other.view(np.ndarray) if isinstance(other, MyArray) else np.asarray(other)
        
        try:
            # 调用自定义equal逻辑
            return self._custom_equal(self_arr, other_arr)
        except Exception as e:
            # 异常直接抛出，不触发默认对象比较
            raise e

    @staticmethod
    def _custom_equal(a, b, **kwargs):
        return np.equal(a, b, **kwargs)

验证效果 #

• 形状匹配时，正常返回布尔数组：

>>> a = MyArray([(1,2,3), (4,5,6)])
>>> b = MyArray([(1,2,3), (4,5,6)])
>>> a == b
array([[ True,  True,  True],
       [ True,  True,  True]])

• 形状不匹配时，直接抛出异常，不再返回False：

>>> c = MyArray([(1,2,3,4,5,6)])
>>> a == c
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 16, in __eq__
  File "<stdin>", line 21, in _custom_equal
numpy.core._exceptions.UFuncOutputCastingError: Cannot cast ufunc 'equal' output from dtype('bool') to dtype('int64') with casting rule 'same_kind'
# 部分NumPy版本会直接抛出形状不匹配异常：
# ValueError: operands could not be broadcast together with shapes (2,3) (1,6)

补充说明 #

• 若需要兼容非数组类型的比较（如标量），可以在__eq__中增加类型判断，避免非数值类型触发异常。
• 保持__array_ufunc__的逻辑，确保其他ufunc操作仍能正常工作。

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