如何指定 Pandas DataFrame 中行的类型？答案

【问题标题】：How can I specify the type of a row in a Pandas DataFrame?如何指定 Pandas DataFrame 中行的类型？
【发布时间】：2023-02-11 06:06:03
【问题描述】：

如果我有一个返回 DataFrame 的函数，并且我知道每一行看起来像某个 Dto 类，那么我应该使用哪个包来指定 DataFrame 的列与 @987654322 的属性相匹配@ 班级？

例如

from dataclasses import dataclass
from dataclasses_json import dataclass_json
from pandas import DataFrame

@dataclass_json
@dataclass
class Dto:
  id: int
  name: str

def get_dataframe() -> DataFrame[Dto]:
  dto: Dto = Dto(id=1, name='alice')
  json: Dict[str, int|str] = dto.to_dict()
  return DataFrame([json])

def use_dataframe(df: DataFrame[Dto]):
  for index, row in df.iterrows():
    # row has type Dto here
    print(f'the id of {row.name} is {row.id}')

在上面的例子中，我想要那些类型提示，这样我就知道我得到的行中的列是什么。我希望能够假设每个 row 都将具有 name 和 id 属性，并且如果我尝试将具有不同列的数据帧发送到该 use_dataframe 函数，则会收到警告。

【问题讨论】：

标签： python pandas dataframe type-hinting

【解决方案1】：

一个选项看起来像dataenforce，上面的解决方案看起来像这样：

from dataenforce import Dataset

DtoDataframe = Dataset["id":int, "name":str]

def get_dataframe() -> DtoDataframe:
  dto: Dto = Dto(id=1, name='alice')
  json: Dict[str, int|str] = dto.to_dict()
  return DataFrame([json])

def use_dataframe(df: DtoDataframe):
  for index, row in df.iterrows():
    # row has type Dto here
    print(f'the id of {row["name"]} is {row["id"]')

不理想，因为我必须使用字符串来指定use_dataframe 中的列。理想情况下，我可以输入row. 并让row.name 和row.id 出现在自动完成中。

【讨论】：