🔌

[FastAPI] さまざまなBaseModelを単一のAPIの入力として受け入れて、入力の種類ごとにバリデーションさせる方法

2022/06/15に公開約1,800字

FastAPIが一つの型しか受け付けない、というお悩み、ありませんか? 解決しますよ。(強引に)

FastAPIで困ったことがありました。
FastAPIでは、バリデーションのために、入力で受け付けるデータ型は固定されてしまっています。
しかし、いろんな型を入力にできたらいいなと思いませんか? 自分はそうでした。
若干強引ですが、コード数をそこまで増やさずに実現する方法を見つけたので紹介します。

大まかな方法としては、
受け入れるBaseModelのクラス名をEnumで保存し、JSONと共にそのクラス名をAPIの入力としてもらい、JSONをクラス名からBaseModelにコンバートするという流れです。

今回は具体例として、2つ以上のBaseModelをAPIの入力として受け入れ、入力に応じて型チェックをする方法を紹介します。
今回は、Model1とModel2の2つをBaseModelを作成し、入力として受け入れることにします。

Step1 BaseModelにclassmethodを追加

いきなり着目する点があるので注意してください。
BaseModelの中に自分自身のクラス名を返すclassmethodを追加します。

1つ目のBaseModel

class Model1(BaseModel):
    #ここに、自分自身のクラス名を呼び出すクラスメソッドを追加しておく。
    @classmethod
    def getClassName(cls):
        return cls.__name__
    #普段通り、ここには型を書く。
    name: str
    val1: str

2つ目のBaseModel

class Model2(BaseModel):
    #ここに、自分自身のクラス名を呼び出すクラスメソッドを追加しておく。
    @classmethod
    def getClassName(cls):
        return cls.__name__
    #普段通り、ここには型を書く。
    name: str
    val2: str
    val3: str

同じ形のクラスメソッドを追加しました。

Step2 EnumにBaseModelを保存

先ほどのクラスたちを、Enumに登録しておきます。ここは人力。

#BaseModelをEnumに保存
class TypeEnum(Enum):
    model1 = Model1.getClassName()
    model2 = Model2.getClassName()

Step3 APIで呼び出す関数の引数にEnumを登録

/test/{model-type}のパスパラメータを先ほど作成したTypeEnum型で指定します。
こうすることで、JSONが何のBaseModelなのかを指定できます。
あとは、それら2つの情報をマッチングさせるだけですね。
入力のデータがバリデーションを通してBaseModelの形になります。

@app.post("/test/{model_type}")
#この引数に着目してください
async def test_post(input_json:str, model_type: TypeEnum):
    #JSONをdictionaryに変換
    input_data = json.loads(input_json)
    
    #dictionary型からmodel_typeのBaseModelに変換する。
    #そのためのコードを文字列から生成する。
    lines = model_type.value+"(**input_data)"
    #文字列をコードとして実行。これが通ればBaseModelのバリデーションも通ったことになる。
    input_data = eval(lines)
    #↑↑↑↑↑↑↑↑これで複数種類ある入力から、今回該当するBaseModel型オブジェクト取得できました
    #あとはよしなにデータをお使いください。今回は、そのままデータを返します。
    return input_data

できましたね。。

Discussion

ログインするとコメントできます