How can I guide FastAPI to pick the right type type from Union

Question

I have the following models:

from pydantic import BaseModel

class A(BaseModel):
   tag: str
   field1: str

class B(BaseModel):
   tag: str
   field1: str
   field2: str

I am using them to define the request body like so:

@app.post('/route')
def handle(request: typing.Union[A, B])

The documentation states:

include the most specific type first, followed by the less specific type

But I was wondering if there was a way I could guide FastAPI to select the correct model based on the tag. Currently, all of my models have a constant tag (i.e. type A always has tag = 'A'). Is there a way I can change my model definition to ensure that the assigned model always matches the tag received in the request body?

Currently, I am getting around this by not typing my handle function, and instead specifically matching the tag to the model, but ideally I would just use typing.Union as the type and be confident it is getting it correct.

try annotating the tag field for A like `tag: Literal['A']` (and similarly for B) — rv.kvetch, Oct 06 '21 at 21:11
I checked answer by @rv.kvetch and it works but there is one small downside. You lose request's autogenerated json body field template in OpenAPI. — devaerial, Oct 07 '21 at 21:39
@rv.kvetch that worked for me, thank you! Do you want to post it as an answer and I can mark it as so? — rbhalla, Oct 13 '21 at 23:19

rbhalla · Accepted Answer · 2022-08-17T10:26:35.453

1

From @rv.kvetch's comment, annotating the tag field to be a string literal, e.g. Literal['A'] meant that correct type was selected every time.

Applying this to the problem above, the code would change to:

class A(BaseModel):
   tag: Literal['A']
   field1: str

class B(BaseModel):
   tag: Literal['B']
   field1: str
   field2: str

edited Aug 17 '22 at 10:26

answered Oct 16 '21 at 19:33

rbhalla

869
8
32

Where exactly should I put the tag? Can you post the full modified code? – Daniel Lavedonio de Lima Aug 01 '22 at 21:57
1

@DanielLavedoniodeLima, I've updated the answer to include it. – rbhalla Aug 17 '22 at 10:26
I realise that for this, given that `field1` & `field2` are required: once u try calling that API without putting `field1` & `field2`, the error message will be a union of that for the missing fields in A & B, which doesn't tell the client what was needed, given the tag – Beast Aug 26 '22 at 08:16
That is true, did you find a way around that @Beast? – rbhalla Sep 05 '22 at 21:54
1

Yes. Refer to class `Item` in `app.py` in [this answer](https://stackoverflow.com/a/71337839/11729954). It uses discriminated unions. To access the data and use it normally like any other Pydantic type, access the `__root__` attribute of `Item.__root__` – Beast Sep 06 '22 at 01:15

How can I guide FastAPI to pick the right type type from Union

1 Answers1