如何为使用自定义数据类作为列名的Pandas DataFrame添加类型提示以通过Mypy检查？

我完全懂你的困扰——用自定义数据类当DataFrame列名在运行时明明一切正常，但mypy却揪着“列名应该是字符串”这点报错。其实只要给DataFrame加上正确的泛型类型提示，就能让mypy识别到列名的真实类型，解决这个问题。

先再梳理下你的场景：你用ColumnName数据类来存储列的元数据（比如总线名、组件名），筛选列时直接访问数据类的属性，运行逻辑没问题，但mypy会抛出"str" has no attribute "bus_name"的错误。

最优解决方案：用Pandas泛型指定列类型 #

从Pandas 1.3.0版本开始，pd.DataFrame支持泛型类型注解，你可以明确告诉mypy：这个DataFrame的列名类型是ColumnName，而不是默认的字符串。

修改后的完整代码如下：

import pandas as pd
from dataclasses import dataclass
from typing import Iterable

@dataclass(frozen=True)  # 建议添加：让数据类不可变，更适合作为列名（Pandas要求列名可哈希且稳定）
class ColumnName:
    bus_name: str
    component_name: str
    component_type: str

    def __hash__(self):
        return hash((self.bus_name, self.component_name, self.component_type))

# 明确标注DataFrame的行索引类型（这里是int）和列名类型（ColumnName）
df_energy: pd.DataFrame[pd.Index[int], ColumnName] = pd.DataFrame({
    ColumnName('bus1', 'comp1', 'type1'): [1, 2, 3],
    ColumnName('bus1', 'comp2', 'type2'): [4, 5, 6],
    ColumnName('bus2', 'comp3', 'type3'): [7, 8, 9],
    ColumnName('bus2', 'comp4', 'type4'): [10, 11, 12],
})

# 现在mypy能正确识别col是ColumnName类型，不会再报错
cols_bus1 = [col for col in df_energy.columns if col.bus_name == 'bus1']

print(df_energy[cols_bus1])

补充说明 #

1. 为什么加frozen=True：数据类作为列名最好是不可变的，frozen=True会自动生成不可变属性，同时也会完善__eq__的实现（你已经自己写了__hash__，保留也没问题，但不可变更符合Pandas对列名的稳定性要求）。

2. 低版本Pandas的兼容方案：如果你的Pandas版本不支持泛型注解，可以在筛选列时用typing.cast强制明确类型，虽然繁琐但能让mypy通过：

   from typing import cast
   
   cols_bus1 = [
       cast(ColumnName, col) 
       for col in df_energy.columns 
       if cast(ColumnName, col).bus_name == 'bus1'
   ]
   

3. 关于Pandera的小补充：你提到用Pandera没成功，确实Pandera更适合固定列结构的场景，而你的列是动态生成的，所以泛型类型提示是更适配的方案。

这样修改后，mypy就能正确识别列名的类型，不会再抛出attr-defined的错误，同时代码的可读性和类型安全性也能得到保障。

备注：内容来源于stack exchange，提问作者Sillywumps