在 Pydantic 中使用 Pandas 数据框作为类型答案

【问题标题】：Using Pandas Data Frame as a Type in Pydantic在 Pydantic 中使用 Pandas 数据框作为类型
【发布时间】：2020-12-22 17:14:11
【问题描述】：

我正在使用pydantic 并希望创建包含熊猫数据框的类。我在网上找了很长时间，没有找到任何东西。我的自定义类型代码如下所示。我将数据帧的类型命名为 pd.DataFrame 但是显然它不正确。有谁知道如何声明 pandas 数据框类型？

import pandas as pd
from pydantic import BaseModel


class SubModelInput(BaseModel):
    a: pd.DataFrame
    b: pd.DataFrame

class ModelInput(BaseModel):
    SubModelInput: SubModelInput
    a: pd.DataFrame
    b: pd.DataFrame
    c: pd.DataFrame

感谢您的帮助！

【问题讨论】：

标签： python pandas dataframe python-typing pydantic

【解决方案1】：

您可以激活Arbitrary Types Allowed:

import pandas as pd
from pydantic import BaseModel


class SubModelInput(BaseModel):
    a: pd.DataFrame
    b: pd.DataFrame

    class Config:
        arbitrary_types_allowed = True

class ModelInput(BaseModel):
    SubModelInput: SubModelInput
    a: pd.DataFrame
    b: pd.DataFrame
    c: pd.DataFrame

    class Config:
        arbitrary_types_allowed = True

【讨论】：

【解决方案2】：

如果我理解正确，您的意图是为pd.Dataframe 创建一个pythonic 类型提示。我想您可以利用以下实现：

import pandas as pd
from pydantic import BaseModel

from typing import TypeVar

PandasDataFrame = TypeVar('pandas.core.frame.DataFrame')


class SubModelInput(BaseModel):
    a: PandasDataFrame
    b: PandasDataFrame


class ModelInput(BaseModel):
    SubModelInput: SubModelInput
    a: PandasDataFrame
    b: PandasDataFrame
    c: PandasDataFrame


data_frame = pd.DataFrame([{"a": "foo", "b": "bar"}])

sub_model = SubModelInput(a=data_frame, b=data_frame)

model = ModelInput(a=data_frame, b=data_frame, c=data_frame, SubModelInput=sub_model)

model.dict()

# {'SubModelInput': {'a':      a    b
# 0  foo  bar, 'b':      a    b
# 0  foo  bar}, 'a':      a    b
# 0  foo  bar, 'b':      a    b
# 0  foo  bar, 'c':      a    b
# 0  foo  bar}

【讨论】：