如何为Pandas DataFrame的DatetimeIndex添加类型提示？

为什么你的Protocol方案失败？ #

Pandas官方类型标注中，pd.DataFrame的index属性声明为pd.Index（或更宽泛的类型），静态类型检查器（如mypy）只会依据声明的类型判断兼容性，不会推断运行时的实际类型。哪怕你创建的DataFrame用了DatetimeIndex，类型检查器仍认为它的index是pd.Index类型，不匹配你定义的TSFrame协议，因此两种调用都会报错。

可行解决方案 #

方案1：子类化pd.DataFrame并显式标注index类型（推荐） #

通过子类化pd.DataFrame，直接将index的类型标注为pd.DatetimeIndex，同时添加运行时断言确保类型正确。这种方式既能保留DataFrame的所有原生方法，又能让类型检查器正确识别索引类型：

import pandas as pd
from typing import override

class TSFrame(pd.DataFrame):
    # 显式声明index类型为DatetimeIndex
    index: pd.DatetimeIndex

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 运行时验证索引类型
        assert isinstance(self.index, pd.DatetimeIndex), "TSFrame requires DatetimeIndex"
    
    # 可选：重写index属性确保类型一致（避免类型检查警告）
    @property
    def index(self) -> pd.DatetimeIndex:
        return super().index  # type: ignore[return-value]
    
    @index.setter
    def index(self, value: pd.DatetimeIndex) -> None:
        super(TSFrame, self.__class__).index.__set__(self, value)

# 使用示例
def test(df: TSFrame):
    # 这里能正确提示DatetimeIndex的专属方法，如to_period
    df.index.to_period("D")

# 合法调用：创建TSFrame实例
tsdf = TSFrame(index=pd.date_range("2022-01-01", "2022-01-02"))
test(tsdf)  # 类型检查不报错

# 非法调用：普通DataFrame无法传入
nontsdf = pd.DataFrame()
test(nontsdf)  # 类型检查器报错

方案2：使用运行时可检查的Protocol #

如果不想子类化，可以给Protocol添加@runtime_checkable装饰器，同时通过类型注解或cast让类型检查器识别符合条件的DataFrame：

import pandas as pd
from typing import Protocol, runtime_checkable

@runtime_checkable
class TSFrame(Protocol):
    index: pd.DatetimeIndex

def test(df: TSFrame):
    df.index.to_period("D")

# 合法调用：显式标注变量类型
tsdf: TSFrame = pd.DataFrame(index=pd.date_range("2022-01-01", "2022-01-02"))  # type: ignore[assignment]
test(tsdf)  # 类型检查不报错

# 非法调用：普通DataFrame
nontsdf = pd.DataFrame()
test(nontsdf)  # 类型检查器报错

注：这里的type: ignore[assignment]是因为Pandas的DataFrame默认类型不匹配TSFrame，需要告诉类型检查器我们确认这个实例符合协议，同时运行时@runtime_checkable会验证索引类型。

方案3：包装DataFrame（严格但繁琐） #

你提到的包装方案虽然繁琐，但能提供最严格的类型隔离，适合需要完全控制接口的场景：

import pandas as pd
from attrs import define, field

@define(auto_attribs=True)
class TSFrame:
    df: pd.DataFrame = field()

    def __attrs_post_init__(self):
        assert isinstance(self.df.index, pd.DatetimeIndex), "Index must be DatetimeIndex"

    @property
    def index(self) -> pd.DatetimeIndex:
        return self.df.index  # type: ignore[return-value]
    
    # 可选：转发DataFrame的其他方法，比如__getitem__
    def __getitem__(self, key):
        return self.df[key]

# 使用示例
def test(df: TSFrame):
    df.index.to_period("D")

tsdf = TSFrame(df=pd.DataFrame(index=pd.date_range("2022-01-01", "2022-01-02")))
test(tsdf)  # 类型检查不报错

nontsdf = TSFrame(df=pd.DataFrame())  # 运行时断言报错

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