Python/NumPy类型标注问题：如何正确声明sensor_data_to_numpy返回值类型

你遇到的问题本质是NumPy异构数组的类型标注，普通同类型数组的标注语法不适用，下面是两种可行方案：

方案1：严谨结构化数组标注（推荐） #

你返回的数组包含固定结构的字段，用NumPy结构化数组可以实现字段级的类型校验：

1. 首先导入依赖类型：

import numpy as np
import numpy.typing as npt
from typing import Union, List
# 此处省略你已定义的SensorData类

2. 定义对应结构的dtype并标注返回值：

# 假设SensorData的matrix字段是固定长度N的浮点序列，按需替换下方的N为实际长度
SensorArrayDType = np.dtype([
    ("label", np.str_),
    ("sample", np.str_),
    ("file", np.str_),
    ("matrix", np.float64, (N,))
])

def sensor_data_to_numpy(data: Union[SensorData, List[SensorData]]) -> npt.NDArray[SensorArrayDType]:
    if isinstance(data, list):
        res = np.empty(len(data), dtype=SensorArrayDType)
        res["label"] = [item.label for item in data]
        res["sample"] = [item.sample for item in data]
        res["file"] = [item.file for item in data]
        res["matrix"] = [item.matrix for item in data]
        return res
    else:
        res = np.empty(1, dtype=SensorArrayDType)
        res["label"][0] = data.label
        res["sample"][0] = data.sample
        res["file"][0] = data.file
        res["matrix"][0] = data.matrix
        return res

如果matrix长度不固定，将dtype中matrix的定义改为("matrix", object)即可，标注逻辑不变。

方案2：快速object类型标注 #

如果你不想修改原有函数的实现逻辑（直接摊平所有元素生成数组），此时NumPy会自动将数组转为object类型，直接标注即可：

def sensor_data_to_numpy(data: Union[SensorData, List[SensorData]]) -> npt.NDArray[np.object_]:
    if isinstance(data, List):
        return np.array([[item.label, item.sample, item.file, *item.matrix] for item in data])
    else:
        return np.array([data.label, data.sample, data.file, *data.matrix])

该方案仅能标注数组的基础类型，无法校验内部元素的结构和具体类型，适合不需要严格类型校验的场景。

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