toml 形式で作ったテーブルの配列を Pydantic でバリデーションして受け取る

背景

toml 形式のファイル

シミュレーションコードの入力部として、toml 形式の入力ファイルを考えている。特に、可変長個のリストを受け渡し可能にして、柔軟なシミュレーションを行うことを想定している。たとえば、以下のような構成である。

[[floor]]
name = "groundfloor"
floor_type = "middle"
ceilingheight = 2.8

    [[floor.room]]
    area = 18.0
    windows = [0, 1, 2, 3]

    [[floor.room]]
    area = 20.0
    windows = [4, 5]

[[floor]]
name = "firstfloor"
floor_type = "middle"
ceilingheight = 2.4

    [[floor.room]]
    area = 35.0
    windows = [6, 7]

[[floor]]
name = "rooftop"
floor_type = "top"
area = 40.0    

floor は任意の階数を積むことができ、各 floor に任意の数の部屋がある。
また、floor のバリエーションとして rooftop という名前の屋上モデルもある。
屋上のモデルは、異なるパラメーターセットを保持している。

TOML のテーブルの配列 を使って記載することで、dict として読み込むことができるが、inp["floor"][0]["room"][0]["area"] などの形式でアクセスすることとなり、カギ括弧が多くなるので、オブジェクトのネスト構造にして、ドット記法でアクセスできるようにすることを考えた。

Pydantic

Pydantic は、"data validation library for Python" である。説明としては、dict ないし json での例示が多い。

このメモでは Pydantic で TOML のテーブルの配列を扱う方法を説明する。これは、dict 内の「dict のリスト」の扱い方と等価である。

Pydantic の version 2.6.1 で試した内容である。

import pydantic
print(pydantic.__version__)
# 2.6.1

コードを作ってみた

ライブラリの読み込みと toml の内容のロード

toml の読み込みは、string io で代替した。

import io
import toml

from pydantic import BaseModel, Field

from typing_extensions import Annotated
from typing import Union, Literal

toml_text = """
[[floor]]
name = "groundfloor"
floor_type = "middle"
ceilingheight = 2.8

    [[floor.room]]
    area = 18.0
    windows = [0, 1, 2, 3]

    [[floor.room]]
    area = 20.0
    windows = [4, 5]

[[floor]]
name = "firstfloor"
floor_type = "middle"
ceilingheight = 2.4

    [[floor.room]]
    area = 35.0
    windows = [6, 7]

[[floor]]
name = "rooftop"
floor_type = "top"
area = 40.0    
"""

with io.StringIO(toml_text) as f:
    inp = toml.load(f)

inp
# {'floor': [{'name': 'groundfloor',
#    'floor_type': 'middle',
#    'ceilingheight': 2.8,
#    'room': [{'area': 18.0, 'windows': [0, 1, 2, 3]},
#     {'area': 20.0, 'windows': [4, 5]}]},
#   {'name': 'firstfloor',
#    'floor_type': 'middle',
#    'ceilingheight': 2.4,
#    'room': [{'area': 35.0, 'windows': [6, 7]}]},
#   {'name': 'rooftop', 'floor_type': 'top', 'area': 40.0}]}

下位の辞書からクラス化してテストする

下位の辞書 room を読めるようにしよう

inp["floor"][0]["room"][0]
# {'area': 18.0, 'windows': [0, 1, 2, 3]}

float と配列を受け取るクラスを Pydantic.Basemodel のサブクラスとして作る
値に対するチェックの例も示している。
極端に狭い部屋は許さない場合、 Field(..., ge=6.0) などとしてここでチェックを実施する。
.model_validate() に適切な dict を渡せばクラスが生成される。

 class Room(BaseModel):
    area: float = Field(..., ge=6.0)
    windows: list[int] = []

myroom = Room.model_validate(inp["floor"][0]["room"][0])
myroom
# Room(area=18.0, windows=[0, 1, 2, 3])

クラス名は Room でなくても動作するが、クラス変数は、テーブルのキーと一致していなければいけない。

Room のリストを含むクラスを作る

toml の中でテーブルの配列に [[floor.room]] という名前を使っているので、これに対応するために、Floor のクラス変数に room を指定して list[Room] というタイプ指定をすれば良い。

class Room(BaseModel):
    area: float = Field(..., ge=6.0)
    windows: list[int]

class Floor(BaseModel):
    name: str
    floor_type: str
    room: list[Room]
    ceilingheight: float

myfloor = Floor.model_validate(inp["floor"][0])
myfloor
# Floor(name='groundfloor', floor_type='middle',
#   room=[Room(area=18.0, windows=[0, 1, 2, 3]), 
#         Room(area=20.0, windows=[4, 5])],
#   ceilingheight=2.8)

別モデルとなる rooftop もクラスを作る

class Rooftop(BaseModel):
    area: float
    floor_type: str

myfloor2 = Rooftop.model_validate(inp["floor"][2])
# Rooftop(area=40.0, floor_type='top')

Floor と Rooftop の混ざったリストを読んで、全体のクラス 'House' を作る

Floor と Rooftop は、クラス変数として floor_type が共通である。
この変数の値をもとに、どちらのクラスを使うかを決める仕組みが、discriminator である。 Annotated を使い、Floor と Rooftop の Union に対して識別の根拠を指定する。この Union をlayer という名前で参照することにした。そして layer のリストを読み込めるように構成する。
RoofTop を混ぜなくても良いなら、もう少しシンプルに floor: list[Floor] とすれば十分である。

class Room(BaseModel):
    area: float = Field(..., ge=6.0)
    windows: list[int]

class Floor(BaseModel):
    name: str
    room: list[Room]
    ceilingheight: float
    floor_type: Literal['middle']

class Rooftop(BaseModel):
    area: float
    floor_type: Literal['top']

Layer = Annotated[Union[Rooftop, Floor], Field(discriminator='floor_type')]

class House(BaseModel):
    floor: list[Layer]

myhouse = House.model_validate(inp)
myhouse
# House(
#     floor=[
#         Floor(name='groundfloor', room=[
#             Room(area=18.0, windows=[0, 1, 2, 3]), 
#             Room(area=20.0, windows=[4, 5])], 
#           ceilingheight=2.8, 
#           floor_type='middle'), 
#
#         Floor(name='firstfloor', room=[
#           Room(area=35.0, windows=[6, 7])], 
#           ceilingheight=2.4, floor_type='middle'), 
# 
#         Rooftop(area=40.0, floor_type='top')
#    ])

以上のコードで、toml から作った dict をネストされたクラスに変換し、myhouse.floor[1].room という形式で内容にアクセスできるように構成できた。

myhouse.floor[1].room
# [Room(area=35.0, windows=[6, 7])]

その他

あとは、必要なチェックを追加したり、パラメータの変換を加えていく。
入力のチェックを充実させるには、Annotated に AfterValidator, BeforeValidator を加えたり、model_validator を使ってひとつのパラメータに閉じない形で、入力のチェックや変換を行うことができる。