将 mypy 与 pandas `to_dict` 方法结合使用

Question

这是我的实际代码的简化版本：

import pandas as pd
from typing import Dict

df = pd.DataFrame(
    {"year": [2024, 2025], "my_output": [1, 2], "foo": ["a", "b"]}
)

df.set_index("year", inplace=True)

row_dict = df.to_dict("index")

ouput_dict: Dict[int, int] = {
    key: val["my_output"] for key, val in row_dict.items()
}

当我对此运行 mypy 时，出现以下错误：

字典理解中的键表达式具有不兼容的类型“Hashable”；预期类型“int”

如何告诉 mypy 我使用“year”列设置的键是整数？

Answer 1

df.to_dict("index")

将数据帧转换为字典，其中键的类型为

Hashable

而不一定是

int

。 mypy 期望

ouput_dict

中的键为

int

。您需要断言键的类型。

import pandas as pd
from typing import Dict, Hashable, Any

df = pd.DataFrame(
    {"year": [2024, 2025], "my_output": [1, 2], "foo": ["a", "b"]}
)

df.set_index("year", inplace=True)

row_dict = df.to_dict("index")

ouput_dict: Dict[int, int] = {
    int(key): val["my_output"] for key, val in row_dict.items() if isinstance(key, int)
}

和

print(output_dict)

{2024: 1, 2025: 2}