Pydantic TypeAdapter でサクサク型付けデシリアライズ - JSON を複数クラスへ自動振り分け
はじめに
最近、仕事の中で Python を触る機会が増え続けています。そんな中、以下のような JSON をデシリアライズするシーンに出会いました。
{
"model_name": "compact",
"weight": 1.5
}
{
"model_name": "powerful",
"turbo": "true"
}
このように、「識別に使えそうな共通のキーはあるが、それ以外のキーが全く異なるような JSON」を読み込んでデシリアライズし、後続で異なる処理を行っていくというものです。
辞書として読み込んだ後に愚直に model_name を参照して条件分岐を書くこともできますが、せっかくであれば型やクラスを活用していきたいですね?(押し付け)
そこで本記事では以下のようなゴールを設定します。
# これは CompactMachine クラスになる
machine = 何かの処理({"model_name": "compact", "weight": 1.5})
# これは PowerfulMachine クラスになる
machine = 何かの処理({"model_name": "powerful", "turbo": "true"})
さあ、これを実現する「何かの処理」はいったい何があるのでしょうか?
前提条件
- Python 3.10 以上(コード例で
CompactMachine | PowerfulMachineのような PEP 604 形式を使うため) - Pydantic v2 以上
本記事の対象読者
- Python で JSON をデシリアライズした結果を適切なクラスに動的にキャストしたい人
- Pydantic が好きな人
Pydantic の discriminator
筆者の以前の記事では、次のような データの一部が異なる JSON を Pydantic を使って型安全に扱うことをトライしました。
// drink_type によって、drink が持つフィールドが異なる。
{
"drink": {
"drink_type": "coffee",
"serve_mode": {
// 省略
}
},
"cup_type": "paper_cup"
}
{
"drink": {
"drink_type": "green_tea",
"region": "famous_region"
},
"cup_type": "my_cup"
}
このケースでは Pydantic の Field で discriminator を活用することができました。
class Coffee(BaseModel):
drink_type: Literal["coffee"] = "coffee"
serve_mode: # 省略
class GreenTea(BaseModel):
drink_type: Literal["green_tea"] = "green_tea"
region: Literal["famous_region", "other_region"]
class ServeRequest(BaseModel):
drink: Coffee | GreenTea = Field(discriminator="drink_type")
cup_type: Literal["paper_cup", "my_cup"]
# ServeRequest.drink は Coffee クラス
req = ServeRequest.model_validate({
"drink": {
"drink_type": "coffee",
"serve_mode": # 省略
},
"cup_type": "paper_cup",
})
# ServeRequest.drink は GreenTea クラス
req = ServeRequest.model_validate({
"drink": {
"drink_type": "green_tea",
"region": # 省略
},
"cup_type": "my_cup",
})
こちらの詳細を知りたい場合は、ぜひ下記の記事もご覧ください。
しかし今回のケースは クラスのプロパティの一部を異なるクラスとしてデシリアライズしたいのではなく、丸ごと異なるクラスとして扱いたい ため異なるケースです。ServeRequest のように「1つのモデルの中に条件分岐を抱え込む」構成ではなく、入力値そのものを CompactMachine や PowerfulMachine といった別クラスのインスタンスに変換したかったので、他のアプローチを探す必要があると分かりました。
Pydantic の TypeAdapter
理想の手段を求め、筆者は Pydantic の森を彷徨いました。。。
そしてついに見つけたのです、TypeAdapter を…!
TypeAdapter を使うことで、下記のような処理が実現できます。
from typing import Annotated, Literal, Union
from pydantic import BaseModel, Field, TypeAdapter
class CompactMachine(BaseModel):
model_name: Literal["compact"]
weight: float
class PowerfulMachine(BaseModel):
model_name: Literal["powerful"]
turbo: Literal["true", "false"]
MachineAdapter = TypeAdapter(CompactMachine | PowerfulMachine)
compact_json = {
"model_name": "compact",
"weight": 1.5,
}
powerful_json = {
"model_name": "powerful",
"turbo": "true",
}
compact = MachineAdapter.validate_python(compact_json)
powerful = MachineAdapter.validate_python(powerful_json)
print(type(compact)) # CompactMachine
print(type(powerful)) # PowerfulMachine
簡潔だ…!!
TypeAdapter は BaseModel.model_validate のように単一モデルをバリデーションするのではなく、Union など任意の型ヒントを事前にコンパイルしておき、その定義に沿って値を振り分けてくれるユーティリティです。validate_python は辞書を渡したときに最適なモデル (CompactMachine / PowerfulMachine) を返してくれるので、以降の処理では Python のクラスとしてメソッドや補完をそのまま使えます。TypeAdapter を使わない場合は自分で model_name を参照して CompactMachine.model_validate を呼び分ける必要があるため、型安全さも記述量も TypeAdapter に軍配が上がります。
playground: https://pydantic.run/store/ba10a1ade69d2c98
discriminator を明示すると安心
ちなみに、先ほどのように何を識別に使うか明示しなくても上手くいく例もありますが、明示しておくにこしたことはありません。
Pydantic 公式ドキュメントにも
"In general, we recommend using discriminated unions. They are both more performant and more predictable than untagged unions, as they allow you to control which member of the union to validate against." — https://docs.pydantic.dev/latest/concepts/unions/
と書かれており、discriminator を挟まない場合は Union の各モデルが「順番に試される」挙動になるため、CompactMachine と PowerfulMachine で似たフィールド構成を持っていると誤って先に並べた方のモデルが選ばれてしまう恐れがあります。将来的に JSON のキーが増えたときの変換事故を防ぐ意味でも、Field(discriminator=...) を使って判別キーを固定しておくのが安全です。
その場合は typing.Annotated と pydantic.Field を併用して以下のように書けます。
# model_name で判別する Union 型
Machine = Annotated[
CompactMachine | PowerfulMachine,
Field(discriminator="model_name"),
]
MachineAdapter = TypeAdapter(Machine)
つまり、冒頭の 何かの処理 は TypeAdapter.validate_python だったんですね!ありがとう Pydantic…!
おわりに
Pydantic を使うことで、デシリアライズした結果を異なるクラスとして扱うことが非常にシンプルに実現できました。
今後もこうしたテクニックを身につけて、型安全で楽しい Python ライフを送りたいと思います。
Discussion