为NumPy/Pandas对象添加Python Typing类型标注的正确方法

我使用NumPy与Pandas已有多年，近期开始用mypy对所有代码进行类型检查。以下示例代码功能正常，但无法通过mypy检查，请问该如何为这段代码添加正确的类型标注？

import pandas as pd
import numpy as np

def get_repeats(date: pd.Timestamp)->np.ndarray:
    return np.repeat(date,3)

if __name__ == '__main__':
    today=pd.Timestamp.utcnow()
    repeats=get_repeats(today)
    print(repeats)

mypy报错信息 #

minimal_example.py:5: error: Argument 1 to "repeat" has incompatible type "Timestamp"; expected "Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]"

解决方案 #

原因 #

mypy报错是因为np.repeat的官方类型注解并未包含pd.Timestamp，虽然运行时Pandas对象能被NumPy兼容处理，但静态类型检查器无法自动识别这种隐式兼容。

修复方式 #

方式1：显式转换为NumPy兼容类型

将pd.Timestamp转换为NumPy的datetime64类型，既符合mypy的类型要求，也让代码逻辑更明确：

import pandas as pd
import numpy as np
from numpy.typing import NDArray

def get_repeats(date: pd.Timestamp) -> NDArray[np.datetime64]:
    # 显式转换为numpy datetime64
    return np.repeat(date.to_numpy(), 3)

这里返回类型使用NDArray[np.datetime64]替代泛泛的np.ndarray，能提供更精确的类型信息，帮助mypy检查后续代码的类型安全性。

方式2：使用类型断言绕过检查

如果不想修改运行逻辑，可以用类型注释告诉mypy忽略当前的参数类型错误：

import pandas as pd
import numpy as np
from numpy.typing import NDArray

def get_repeats(date: pd.Timestamp) -> NDArray[np.datetime64]:
    return np.repeat(date, 3)  # type: ignore[arg-type]

这种方式适合快速解决检查报错，但缺乏严谨性，建议优先选择显式转换的方案。

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