Using Pandas Data Frame as a Type in Pydantic

Question

I'm using pydantic and want to create classes which contain pandas dataframes. I was looking for this online for quite a time and did not find anything. My code for the custom types looks as following. I named the type for dataframes pd.DataFrame but obviously its not correct. Does anyone know how to declare a pandas dataframe type?

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

Thanks for any help!

score 7 · Answer 1 · answered Jan 07 '21 at 15:00

If I understand correctly, your intention is to create a pythonic type hint for a pd.Dataframe. I suppose you could utilize the below implementation:

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}

score 7 · Accepted Answer · answered May 20 '21 at 07:39

You can activate 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

score 1 · Answer 3 · answered Jul 27 '23 at 13:21

For pydantic V2.0, there is a slight change in the config. The right way to do it is:

import pandas as pd
from pydantic import BaseModel, ConfigDict


class SubModelInput(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    a: pd.DataFrame
    b: pd.DataFrame

class ModelInput(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    SubModelInput: SubModelInput
    a: pd.DataFrame
    b: pd.DataFrame
    c: pd.DataFrame

More info can be found in the docs

Using Pandas Data Frame as a Type in Pydantic

3 Answers3