为含Drake表达式的NumPy数组添加正确类型提示的方法

问题背景 #

当创建包含pydrake.symbolic.Expression或Formula等Drake符号对象的NumPy数组时，直接使用npt.NDArray[pydrake.symbolic.Expression]会触发类型检查错误，错误信息如下：

Could not specialize type "NDArray[ScalarType@NDArray]"
Type "Expression" cannot be assigned to type "generic"
"Expression" is incompatible with "generic"

用户的初始代码示例：

import numpy as np
import numpy.typing as npt
import pydrake.symbolic

def some_function(expr: pydrake.symbolic.expression) -> npt.NDArray[pydrake.symbolic.expression]:
   return np.array([expr])

解决方案 #

NumPy的NDArray类型提示要求标量类型属于其认可的数值类型体系，而Drake的符号对象不属于此类。要明确标注数组内容类型，可采用以下两种方案：

方案1：用np.dtype包装Drake类型 #

通过np.dtype[pydrake.symbolic.Expression]指定数组元素类型，同时创建数组时显式声明dtype=object（避免NumPy自动尝试转换类型）：

import numpy as np
import numpy.typing as npt
import pydrake.symbolic

def some_function(expr: pydrake.symbolic.Expression) -> npt.NDArray[np.dtype[pydrake.symbolic.Expression]]:
    return np.array([expr], dtype=object)

方案2：使用typing.Sequence（非严格NDArray场景） #

如果不需要严格标注为NumPy数组，仅需明确元素类型，可使用Sequence结合具体Drake类型：

from typing import Sequence
import numpy as np
import pydrake.symbolic

def some_function(expr: pydrake.symbolic.Expression) -> Sequence[pydrake.symbolic.Expression]:
    return np.array([expr])

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