使用 Pydantic 和 Pandera 初始化 Class 属性

Question

我是 Pydantic 和 Pandera 的新手，需要 class 实例化和初始化方面的帮助。

我在一个文件sim.py中有以下代码：

import pandera as pa
from pydantic import BaseModel
from datetime import datetime

class ScheduleDF(pa.SchemaModel):
    person_id: Series[int] = pa.Field(ge=0, coerce=True)
    shift_id: Series[int] = pa.Field(ge=0, coerce=True)
    start_time: Series[datetime]
    end_time: Series[datetime]

class Schedule(BaseModel):
    schedule_df: DataFrame[ScheduleDF]
    events_df: DataFrame[EventsDF]

    @pa.check_types
    def initialize_from_df(self, schedule_df: DataFrame[ScheduleDF]):
        self.schedule_df = schedule_df

以及另一个文件sim_test.py中的以下代码：

from sim import ScheduleDF, Schedule

def test_schedule():
    y = 2022
    m = 9
    d = 1

    schedule_df = DataFrame[ScheduleDF](
        {'person_id': [1, 2], 'shift_id': [10, 20],
         'start_time': [datetime(y, m, d, 0, 0, 1), datetime(y, m, d, 0, 0, 5)],
         'end_time': [datetime(y, m, d, 0, 0, 3), datetime(y, m, d, 0, 0, 6)]
        }
    )
    sample_schedule = Schedule()
    sample_schedule.initialize_from_df(schedule_df)

test_schedule()

当我运行sim_testing.py时，我收到以下错误：

pydantic.error_wrappers.ValidationError: 2 validation errors for Schedule
schedule_df
field required (type=value_error.missing)
events_df
field required (type=value_error.missing)

我可以看到events_df丢失的原因 - 我没有在test_schedule()内初始化它。 但是，似乎我已经初始化了schedule_df 。

我尝试在@pa.check_types装饰器上方为initialize_from_df()添加@classmethod并按照此处和此处的建议将 function 中的self更改为cls ，但它仍然给了我同样的错误。 这似乎是一个 Pydantic 问题，而不是 Pandera 问题。

我会很感激一些帮助弄清楚发生了什么以及如何纠正它。 谢谢！

Answer 1

在您的Schedule class 中，您定义了两个必填字段schedule_df和event_df 。

class Schedule(BaseModel):
    schedule_df: DataFrame[ScheduleDF]
    events_df: DataFrame[EventsDF]

因此，当您尝试使用sample_schedule = Schedule()实例化它时，它必然会失败，因为您没有为这两个必填字段提供任何值。

基本上，您有两种前进方式：

1. 像这样在实例化时传递ScheduleDF和EventsDF的实例（注意：你忘了给我们EventsDF class 的定义，所以我只是编一个）：

schedule_df = DataFrame[ScheduleDF](
    {
        'person_id': [1, 2],
        'shift_id': [10, 20],
        'start_time': [datetime(y, m, d, 0, 0, 1), datetime(y, m, d, 0, 0, 5)],
        'end_time': [datetime(y, m, d, 0, 0, 3), datetime(y, m, d, 0, 0, 6)]
    }
)

events_df = DataFrame[EventsDF](
    {
        'event_id': [1],
        'start_time': [datetime(y, m, d, 0, 0, 3)],
        'end_time': [datetime(y, m, d, 0, 0, 4)]
    }
)

sample_schedule = Schedule(schedule_df=schedule_df, events_df=events_df)

2. 或者，如果您以后真的想为schedule_df和events_df分配值，请在 class 定义中将它们设为可选：

from typing import Optional


class Schedule(BaseModel):
    schedule_df: Optional[DataFrame[ScheduleDF]]
    events_df: Optional[DataFrame[EventsDF]]

这样，调用sample_schedule = Schedule()就可以了，而sample_schedule基本上将包含schedule_df=None events_df=None 。