Zenn
📌

PydanticのRootModelを使いこなす

2023/09/22に公開
Python
FastAPI
Pydantic
tech

PydanticのBaseModelの恩恵をlistやdictにも与えたい

こんにちは、極論モンスターのYosematです。

PydanticはTypeSafeにPythonをかけて最高ですよね。特に私が好きなのはmodel_validateのようなSerialize/Deserializeの機能です。しかしPydanticのBaseModelはうまいことdictやjsonへのSerialize/Deserializeができますがlist[MyModel]はPydanticのクラスではないのでmodel_validateのようなクラスメソッドにアクセスできません。

ダサい方法を先に書きます。

from __future__ import annotations
from pydantic import BaseModel

class MyModel(BaseModel):
    value: str


class MyModels(BaseModel):
    values: list[MyModel]

    @classmethod
    def custom_validate(cls, obj: list[dict[str, str]]) -> MyModels:
        return cls(values=[MyModel.model_validate(item) for item in obj])

でも自前でcustom_validateを実装するのってPydanticを使う意義が薄れる感じがして悔しいですよね。

RootModel

そんなときRootModelです。RootModelはlistやdictなどのコンテナ型をBaseModelへと変換するWrapperです。Generic引数として受け付けたい型を書いてやります。今回はMyModelのlistを定義します。中身には.rootでアクセスできます。

from pydantic import BaseModel, RootModel

class MyModel(BaseModel):
    value: str


MyModels = RootModel[list[MyModel]]


def main():
    obj = [{"value": "hoge"}, {"value": "fuga"}]
    mymodels = MyModels.model_validate(obj)
    print(mymodels.root)  # [MyModel(value='hoge'), MyModel(value='fuga')]

最高でしょう？

RootModel

上記の方法は軽量でオススメですが、MyModelsクラスにメソッドがはやせません。RootModelを継承するアプローチを利用するとこの問題を解決できます。

from pydantic import BaseModel, RootModel


class MyModel(BaseModel):
    value: str


class MyModels(RootModel[list[MyModel]]):
    def custom_method(self):
        print(f"This is a custom method with root: {self.root}")


def main():
    obj = [{"value": "hoge"}, {"value": "fuga"}]
    mymodels = MyModels.model_validate(obj)
    mymodels.custom_method()  # This is a custom method with root: [MyModel(value='hoge'), MyModel(value='fuga')]

まとめ

最高でしょう？（押しつけ）
ためになった方はいいね押して行っていただけますと幸いです。

参考

Pydantic公式ドキュメント

Discussion