Pydantic モデルのクラス変数に参照して加工した値を定義したい

Pydantic使っていた時に行き当たったちょっとした悩みをネットの海に放っておくといった趣旨の記事です。

やりたかったこと

やりたいことは以下でした。（うまい例は思い浮かばなかった）

商品名(name)、金額(amount)、保存先ファイル名(destination)をjson形式でファイルとログに出力する
商品名(name)、金額(amount)はリクエストで受け取ったものを入れる
ファイル名は商品名_日付を入れる（e.g. keyboard_20210211）

そこでPydanticで出力したい値をクラス変数にもったBaseModel: Purchaseクラスを作っていました。

悩んだこととしては、ファイル名をどうやって持たせるかでした。

最悪model.dict()で出力した後に追加するを考えていましたが、せっかくModelを用意したのにこれでは綺麗じゃないから避けたかったです。

import datetime

from typing import Any, Dict, Optional

from pydantic import BaseModel, Field


class Purchase(BaseModel):
    name: str = Field(..., min_length=1)
    amount: int = Field(..., ge=0)
    destination: Optional[int] = None

    def filename(self) -> str:
        d = datetime.date.today()
        return self.name + '_' + d.strftime('%Y%m%d')


def main() -> Dict[str, Any]:
    # リクエストボディの値を渡す
    purchase = Purchase(
        name="keyboard",
        amount=5000,
    )

    # purchaseインスタンスを生成するタイミングでdestinationも入って欲しい
    output = purchase.dict()
    output['destination'] = purchase.filename()

    return output


if __name__ == '__main__':
    print(main())

実現方法

root_validatorsを使えばやりたいことできるかもよと教えてもらいました。

pydanticのvalidator

インスタンスが生成されるタイミングでデコレータで指定した変数のチェックを行える。

class Purchase(BaseModel):
    name: str = Field(..., min_length=1)
    amount: int = Field(..., ge=0)
    destination: Optional[int] = None

    @validator('*')
    def hoge(cls, v):
        print(f'hoge: {v}')
        return v


# 実行結果
# hoge: keyboard
# hoge: 5000
# {'name': 'keyboard', 'amount': 5000, 'destination': 'keyboard_20210211'}

root_validator

root_validatorを使うと全クラス変数を参照できるようになります。

class Purchase(BaseModel):
    name: str = Field(..., min_length=1)
    amount: int = Field(..., ge=0)
    destination: Optional[int] = None

    @root_validator()
    def fuga(cls, values):
        print(f'fuga: {values}')
        return values

# 実行結果
# fuga: {'name': 'keyboard', 'amount': 5000, 'destination': None}
# {'name': 'keyboard', 'amount': 5000, 'destination': 'keyboard_20210211'}

今回やったこと

filename()をroot_validatorの中で実装させるようにしました。

import datetime

from typing import Any, Dict, Optional

from pydantic import BaseModel, Field, root_validator


class Purchase(BaseModel):
    name: str = Field(..., min_length=1)
    amount: int = Field(..., ge=0)
    destination: Optional[int] = None

    @root_validator()　# 追記
    def filename(cls, values: dict) -> Dict[Any, Any]:
        d = datetime.date.today()
        values['destination'] = values['name'] + '_' + d.strftime('%Y%m%d')
        return values


def main() -> Dict[str, Any]:
    # リクエストボディの値を渡す
    purchase = Purchase(
        name="keyboard",
        amount=5000,
    )

    return purchase.dict()


if __name__ == '__main__':
    print(main())

おわりに

ひとまず微妙な書き方をせずには済みました。

ただ検証を行うところでこういうのをやってしまってよかったのか🤔

日々のあれこれφ(．．)

もっぱら壁打ち