pydanticでUnion型を使うときはLiteralと組み合わせるといい

3519 ワード
typing Python typing-extensions pydantic Python テキストリンク

チャットボットを実装しており、「ユーザーの過去のリクエスト内容をRedisに保存しておいて、必要に応じて参照してbotの返信に利用する」という実装をしています。具体的には以下のような仕様とプログラムです。

  1. ユーザーのリクエストのjsonをRequest型にキャストして利用する
  2. Redisでは、pydanticの型をjsonに変換して文字列として保存し、取り出すときにRequest型にキャストする
  3. 同時にそのjsonをログに出力する
from __future__ import annotations
from typing import Union, Optional
from typing_extensions import Literal
from .event import EventData


class Request(BaseModel):
    user_id: str
    session_id: str
    command: Union[MessageCommand, ButtonCommand]


class MessageCommand(BaseModel):
    """自由文のメッセージが入力されたとき"""
    type: Literal["message"]
    message: Optional[str]
    is_first: bool = False


class ButtonCommand(BaseModel):
    """ボタンを入力されたとき"""
    type: Literal["button"]
    event: Optional[EventData]
    is_first: bool = False


Request.update_forward_refs()

当初、typeのプロパティを用意しておらず、Redisから取り出す際にButtonCommandであるべきjsonが、誤って空のMessageCommandとしてキャストされてしまっていました。デフォルト値やOptionalを利用しているときは、間違ったjsonでもキャストできてしまいます。

それをtypeプロパティにLiteralで固定の文字列を指定することで、確実に変換先を指定することができます。

参考