Les Crash-Devs d'un Passionné

Union discriminant avec pydantic

/Catégorie/python

Temps de lecture : 4 minutes

Un objet

pydantic
peut avoir de nombreux attributs et notamment avoir le bon type en fonction d'une valeur

Dans l'exemple ci-dessous on veut associer un

User
à un
Customer

clipboard
Copier le code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from pydantic import BaseModel

class MinorUser(BaseModel):
    age: int
    name: str

class MajorUser(BaseModel):
    age: int
    name: str

class Customer(BaseModel):
    id: int
    user: MinorUser | MajorUser

data = {
    "id": 12,
    "user": {
        "age": 40,
        "name": "David"
    }
}

customer = Customer(**data)
print(customer) # id=12 user=MinorUser(age=40, name='David')

On remarque que user est de type

MinorUser
. Si toutefois je change l'ordre de déclaration par
user: MajorUser | MinorUser
j'obtiens maintenant un type
MajourUser
à la place

id=12 user=MajorUser(age=40, name='David')

Ce comportement n'est pas souhaitable, on cherche à obtenir le bon type en fonction d'un attribut particulier, ici l'âge.

On va ainsi préciser à

pydantic
de regarder la valeur de
age
pour choisir le bon type. Voici l'exemple complet ci-dessous

clipboard
Copier le code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from typing import Annotated, List, Literal
from pydantic import BaseModel, Field, field_validator

class MinorUser(BaseModel):
    status: Literal["minor"]
    age: int
    name: str

class MajorUser(BaseModel):
    status: Literal["major"]
    age: int
    name: str

class Customer(BaseModel):
    id: int
    user: Annotated[MinorUser | MajorUser, Field(discriminator="status")]


    @field_validator("user", mode="before")
    @classmethod
    def validate(cls, user_data: User, validation_info) -> User:
        user_data["status"] = "major" if user_data["age"] >= 18 else "minor"
        return user_data

data = {
    "id": 12,
    "user": {
        "age": 15,
        "name": "David"
    }
}

customer = Customer(**data)
print(customer) # id=12 user=[MinorUser(status='minor', age=15, name='David')]

On remarque plusieurs points intéressants :

- On ajoute un

status
dans
MinorUser
et
MajorUser
qui sera le champ que le modèle
Customer
va inspecter. En fonction de sa valeur, il va choisir
MajorUser
ou
MinorUser
. Il faut donc, dans ces derniers, mettre la valeur souhaitée dans le champ
status
. Il est important de noter qu'il faut mettre un
Literal
car vis-à-vis de
Customer
c'est une liste qui est proposée, il doit choisir entre
MajorUser
ou
MinorUser

- Sur les objets

MinorUser
et
MajorUser
j'ai mis un
Literal
sur status car il faut bien utiliser le typage d'une liste (et non la structure
list
directement) avec le mot-clé
Literal
pour que
pydantic
puisse s'y retrouver

- Sur

Customer
il faut préciser le champs discriminant et la liste des types disponible

user: Annotated[MinorUser | MajorUser, Field(discriminator="status")]

Retrouver l'explication de

Annotated
sur l'article dédié : Annotated avec pydantic

Le champ

Field
provient de
pydantic
et permet d'utiliser le discriminant

- L'utilisation d'un discriminant doit être obligatoirement fait en mode

before
, c'est-à-dire avant que
pydantic
lance la validation. Je l'ai mis sur la validation du champ
user
où je mets le check de l'
age
pour affecter le
status
dynamiquement en fonction de la valeur de l'âge

clipboard
Copier le code
1
2
3
4
5
@field_validator("user", mode="before")
@classmethod
    def validate(cls, users_data: User, validation_info) -> User:
        users_data["status"] = "major" if users_data["age"] >= 18 else "minor"
        return users_data

Maintenant si on injecte ce payload

clipboard
Copier le code
1
2
3
4
5
6
7
data = {
    "id": 12,
    "user": {
        "age": 15,
        "name": "David"
    }
}

On obtient

id=12 user=[MinorUser(status='minor', age=15, name='David')]

Et avec ce payload

clipboard
Copier le code
1
2
3
4
5
6
7
data = {
    "id": 12,
    "user": {
        "age": 40,
        "name": "David"
    }
}

On obtient

id=12 user=[MajorUser(status='minor', age=40, name='David')]

De cette manière on a un objet qui obtient le bon type en fonction d'une règle métier