如何为字符串dtype编写__array_ufunc__ 适配NDArrayOperatorsMixin

你遇到的报错核心是NDArrayOperatorsMixin默认的__array_ufunc__实现会强制将ufunc的输出转为你的自定义类类型，但字符串相等这类运算的输出是bool dtype，你原本的实现既没有兼容非数值dtype的构造逻辑，也没有处理部分ufunc输出不需要包裹为自定义类的场景。

具体修改步骤 #

• 去掉自定义类构造函数中对输入数据的数值dtype校验逻辑，允许字符串、bool等非数值类型传入
• 重写__array_ufunc__方法，增加输出类型判断逻辑：如果ufunc运算结果的dtype和输入自定义类的存储dtype不匹配（比如比较类操作输出bool），可以直接返回numpy原生数组，也可以根据业务需要包裹为自定义类
• 完善ufunc的输入输出调度逻辑，避免自定义类实例干扰numpy原生ufunc的类型匹配

完整可运行的修改后代码示例 #

import numpy as np
from numpy.lib.mixins import NDArrayOperatorsMixin

class MyNewArrayClass(NDArrayOperatorsMixin):
    def __init__(self, values):
        # 移除原有的数值dtype校验，支持所有合法numpy dtype
        self._data = np.asarray(values)
    
    def __array__(self, dtype=None):
        return np.asarray(self._data, dtype=dtype)
    
    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        # 把所有输入里的自定义类实例转成原生numpy数组
        input_args = []
        for inp in inputs:
            input_args.append(inp._data if isinstance(inp, MyNewArrayClass) else inp)
        
        # 处理输出参数里的自定义类实例
        if 'out' in kwargs:
            out_args = []
            for out in kwargs['out']:
                out_args.append(out._data if isinstance(out, MyNewArrayClass) else out)
            kwargs['out'] = tuple(out_args)
        
        # 执行原生numpy ufunc运算
        result = getattr(ufunc, method)(*input_args, **kwargs)
        
        # 处理返回值：只有dtype匹配时才包裹为自定义类，否则直接返回原生数组
        if isinstance(result, tuple):
            return tuple(MyNewArrayClass(res) if res.dtype == self._data.dtype else res for res in result)
        else:
            return MyNewArrayClass(result) if result.dtype == self._data.dtype else result

测试效果 #

运行你给出的测试代码不会再报错：

x = MyNewArrayClass(np.array(["a", "b", "c"]))
print(x == "a")
# 输出：[ True False False]

补充说明 #

如果你业务要求比较类操作的结果也要包裹为自定义类，只需要把返回逻辑里的类型判断去掉，所有结果都用MyNewArrayClass包裹即可，只要构造函数没有dtype限制就能正常运行。

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