Search code examples
pythonjsonpydantic

Define constraints for different JSON schema models in a list


I have some JSON with a structure similar to what is shown below. The threshold list represents objects where the type can be "type": "upper_limit" or "type": "range". Notice the "target" value should be an integer or float depending on the type of the object.

{
    "name": "blah",
    "project": "blah blah",
    "threshold": [
        {
            "id": "234asdflkj",
            "group": "walkers",
            "type": "upper_limit",
            "target": 20,
            "var": "distance"
        },
        {
            "id": "asdf34asf2654",
            "group": "runners",
            "type": "range",
            "target": 1.7,
            "var": "speed"
        }
    ]
}

Pydantic models to generate a JSON schema for the above data are given below:

class ThresholdType(str, Enum):
    upper_limit = "upper_limit"
    range = "range"


class ThresholdUpperLimit(BaseModel):
    id: str
    group: str
    type: ThresholdType = "upper_limit"
    target: int = Field(gt=2, le=20)
    var: str


class ThresholdRange(BaseModel):
    id: str
    group: str
    type: ThresholdType = "range"
    target: float = Field(gt=0, lt=10)
    var: str


class Campaign(BaseModel):
    name: str
    project: str
    threshold: list[ThresholdUpperLimit | ThresholdRange]

The models validate the JSON, but the constraints for the target value are being ignored for the type. For example, if a threshold object contains "type": "range", "target": 12, then no errors are thrown because it is being parsed as an integer and therefore constraints defined by the ThresholdUpperLimit are used; but the constraints defined by ThresholdRange should be used because the type is "range". Any suggestions on how to properly handle this?


Solution

  • I managed to enforce the correct model resolution by changing your use of enum subclassing to using a Literal.

    from pydantic import BaseModel, Field
    from typing import Literal
    
    
    class ThresholdUpperLimit(BaseModel):
        id: str
        group: str
        type: Literal["upper_limit"]
        target: int = Field(gt=2, le=20)
        var: str
    
    
    class ThresholdRange(BaseModel):
        id: str
        group: str
        type: Literal["range"]
        target: float = Field(gt=0, lt=10)
        var: str
    
    
    class Campaign(BaseModel):
        name: str
        project: str
        threshold: list[ThresholdUpperLimit | ThresholdRange]
    

    However, this does not force target to be a float on the ThresholdRange class. An int will pass.