Zenn
🙆

今のところの自分好みのPydantic v2のBaseModel

2024/12/07に公開
Python
Pydantic
tech

自分用のBaseModel

pydanticはPythonで構造体のようなものを作るときにとても便利だ。

しかし、デフォルトではいくつか自分好みでないところがある。

そこで、pydanticを使うときは自分用のBaseModelを作成して使い回すことにした。

from io import StringIO
from typing import Any

import pydantic
import polars as pl
from pydantic.alias_generators import to_camel
from pydantic import model_validator, field_validator


class BaseModel(pydantic.BaseModel):
    model_config = pydantic.ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
        arbitrary_types_allowed=True,
        validate_assignment=True,
        use_enum_values=True,
        json_encoders={
            pl.DataFrame: lambda df: df.serialize(format="json"),
        },
    )

どのような設定か

これで以下のような設定になった。

  • alias_generator=to_camelでフィールド名をlowerCamelCaseにする
  • populate_by_name=Trueで元のフィールド名も使えるようにする
    • JSONでlowerCamelCaseにしたいけど、Pythonではsnake_caseのままにしたい
  • arbitrary_types_allowed=Trueで任意の型を入れられるようにする
    • これがないとpl.DataFrameなどを入れられない
  • validate_assignment=Trueでフィールドに代入を行った時にもValidationされるようにする
    • デフォルトではインスタンス生成時のみしかValidationされない
  • use_enum_values=Trueでenum.Enum値も使えるようにする
  • json_encoders={...}でJSONエンコード時に値を変換するようにする
    • デフォルトでは対応してない型もJSONエンコードしたい場合に入れる

自分用のBaseModelを使用したクラス

また、BaseModelを継承するクラス側では、以下のようにすればJSONからサードパーティの型(e.g. pl.DataFrame)の値を生成できるようになる。

# 自分のBaseModelを継承したクラス
class CustomData(BaseModel):
    li: list[str]
    df: pl.DataFrame

    @field_validator("df", mode="before")
    @classmethod
    def preprocess_df(cls, value: Any) -> pl.DataFrame:
        """validate時に、DataFrameに変換する"""
        if isinstance(value, str):
            df = pl.DataFrame.deserialize(StringIO(value), format="json")
            return df
        else:
            return value
  • field_validatorでCustomData#model_validate_json時にpl.DataFrame#serializeされた値からpl.DataFrameを生成できるようにする

この設定に至るまで

モデル値(Python)　↔︎　JSON

pydantic v2になってからあまり日本語の情報が出回っておらず、どのように自分好みにするのか探し回った。

PythonではPEP8通りsnake_caseでフィールドを使いたいし、JSONはJavaScriptのコーディングスタイルに合わせてlowerCamelCaseにしたい。

型チェック

型ヒント(+ mypy)とpydanticがあれば静的型付け言語のようにできるじゃないかと思ったが、勘違いもあった。

特にpydanticを使っているうちにValidationはいつでも行われていると錯覚していたが、インスタンス生成時のみがデフォルトだった。

例えば、以下のようなクラスとしてもインスタンス生成後はValidationが行われず、値をセットすることができてしまう。

class CustomData2(pydantic.BaseModel):
    li: list[str]

data2 = CustomData2(li=["a", "b", "c"])
data2.li = -1  # OK

そして、自分用のBaseModelを使用してもValidationできないケースもあることを頭の片隅に置いておきたい。

data = CustomData(li=["a", "b", "c"], df=pl.DataFrame())
data.li = -1  # NG
data.li.append(-1)  # OK

data.li = -1はフィールドへの代入なのでpydanticが__setitem__を内部で定義すれば挙動を変えられそうだが、data.li.appendは不可能なところと関連するのだろうと思う。

React.jsの状態の変更にはsetStateを使わなければならないのと同様の制限がpydanticにあるのは不思議ではない。

モデルの保存・復元

モデル値を生成したら、それを保存したり復元したりできるようにしたい。
初めはpickleで行っていたが、エラーが多発したため、後から人が保存したファイルを修正できるようにJSONとして保存したくなった。

pl.DataFrameなど一見してどうJSONにしたら良いかわからないものも、エンコード・デコード時の挙動をカスタマイズすれば保存・復元が可能になった。

その他

その他、model_validatorでいわゆる不変表明を行うことができる。

Discussion