如何使用mypy标注包含自定义numpy dtype的数组？

你遇到的问题是因为npt.NDArray的类型参数需要的是类型注解层面的dtype类型，而非运行时的np.dtype实例。直接传入MY_DTYPE（一个运行时实例）会让mypy无法识别，所以需要用类型别名来关联自定义dtype的类型信息。

正确的实现方式 #

import numpy as np
import numpy.typing as npt
from typing import TypeAlias

# 1. 定义结构化dtype的规格描述
VERSION_DTYPE_SPEC = [("major", "u2"), ("minor", "u2"), ("patch", "u2")]

# 2. 创建运行时使用的dtype实例
VERSION_DTYPE = np.dtype(VERSION_DTYPE_SPEC)

# 3. 定义类型注解专用的类型别名（关联dtype的类型信息）
VersionDType: TypeAlias = np.dtype[VERSION_DTYPE_SPEC]

# 示例数组
my_array = np.array([(1, 2, 3), (4, 5, 6)], dtype=VERSION_DTYPE)

# 函数注解使用类型别名
def print_entries(array_var: npt.NDArray[VersionDType]) -> None:
    """Print array entries."""
    for row in array_var:
        print(row)

关键说明 #

• VersionDType作为类型别名，本质是告诉mypy这个类型对应我们自定义的结构化dtype，而非直接传入运行时的实例。
• 拆分dtype规格、运行时实例、类型别名三个部分，既满足运行时的数组创建需求，又能让静态类型检查工具（如mypy）正确识别类型约束。

如果不想额外定义类型别名，也可以直接在注解中写dtype的类型信息，但代码会比较冗长：

def print_entries(array_var: npt.NDArray[np.dtype[([("major", "u2"), ("minor", "u2"), ("patch", "u2")],)]]) -> None:
    ...

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