本文介绍如何使自定义类（如带单位的 quantity）支持 numpy 通用函数（如 np.sin、np.sqrt），关键在于正确实现 `__array_ufunc__`，并处理标量与序列输入的双重场景，避免 typeerror。
要让自定义类（例如 Quantity）自然融入 NumPy 的数学运算生态——即支持 np.sin(q)、np.sqrt([q1, q2]) 等形式——核心在于正确实现 __array_ufunc__ 方法。但需注意：仅在 __array_ufunc__ 中简单提取 .value 并转发给 ufunc 是不够的，尤其当输入为 Python 列表或元组等非数组容器时，NumPy 会尝试对每个元素调用 ufunc，而此时 Quantity 实例本身不满足 ufunc 的输入要求（缺少 __array__ 或可转换协议），从而抛出 TypeError: loop of ufunc does not support argument 0 of type Quantity...。
正确的做法是：在 __array_ufunc__ 中主动识别并“解包”含 Quantity 实例的容器（如 list、tuple），将其统一转换为纯数值结构（如 np.ndarray 或 list of numbers），再交由 ufunc 处理。以下是一个健壮、可复用的实现：
import numpy as np
class Quantity:
def __init__(self, value, units):
self.value = np.asarray(value) # 支持标量/数组输入
self.units = units
def __array_ufunc__(self, ufunc, method, *args, **kwargs):
if method != "__call__":
return NotImplemented # 仅支持函数调用，不支持 reduce/at/accumulate 等
# 将 args 中所有 Quantity 实例替换为其 .value，其余保持不变
def convert_arg(arg):
if isinstance(arg, Quantity):
return arg.value
elif isinstance(arg, (list, tuple)):
# 递归处理容器：将其中的 Quantity 替换为 .value，其他保留
return type(arg)([
a.value if isinstance(a, Quantity) else a for a in arg
])
else:
return arg
new_args = tuple(convert_arg(a) for a in args)
return ufunc(*new_args, **kwargs)
def __repr__(self):
return f"Quantity({self.value}, '{self.units}')"✅ 效果验证：
q = Quantity(0.5, "rad") qs = [Quantity(0.0, "rad"), Quantity(np.pi/2, "rad")] print(np.sin(q)) # → sin(0.5) ≈ 0.4794 print(np.sin(qs)) # → array([0., 1.]) print(np.add(q, q)) # → 1.0 print(np.maximum(qs[0], qs[1])) # → π/2
⚠️ 重要注意事项：
立即学习“Python免费学习笔记（深入）”；
- 不要实现 __array__（除非你明确希望该类被无条件视为 np.ndarray 子集），否则会意外启用大量 NumPy 函数（包括你不希望支持的操作，如 np.concatenate、np.reshape），破坏语义封装；
- __array_ufunc__ 仅接管 ufunc（如 np.sin, np.add），不影响 np.array() 构造或索引等行为；
- 若需支持 np.array([q1, q2]) 返回带单位的数组，应额外实现 __array__ 或 __array_function__（用于覆盖 np.array 等非 ufunc 函数）；
- 对于高阶操作（如 np.vectorize、np.frompyfunc），仍需单独适配，__array_ufunc__ 不生效。
综上，通过精细化控制 __array_ufunc__ 中的参数转换逻辑（尤其是容器类型解包），即可以统一、声明式的方式支持全部 NumPy 通用函数，无需为每个函数单独定义方法，兼顾简洁性与健壮性。
