'Pydantic child discrimination from parent field

I am trying to create a limited set of Pydantic models from a boto3 response (specifically CodeDeploy's get_deployment). Where I'm having trouble is this:

        'revision': {
            'revisionType': 'S3'|'GitHub'|'String'|'AppSpecContent',
            's3Location': {
                'bucket': 'string',
                'key': 'string',
                'bundleType': 'tar'|'tgz'|'zip'|'YAML'|'JSON',
                'version': 'string',
                'eTag': 'string'
            },
            'gitHubLocation': {
                'repository': 'string',
                'commitId': 'string'
            },
            'string': {
                'content': 'string',
                'sha256': 'string'
            },
            'appSpecContent': {
                'content': 'string',
                'sha256': 'string'
            }
        }

I have tried a discriminated union, but I'm missing something

class GitHubRevision(BaseModel):
    revisionType: Literal['GitHub']
    repository: str
    commitId: str

class StringRevision(BaseModel):
    revisionType: Literal['String']

class RevisionModel(BaseModel):
    revision: Union[StringRevision, GitHubRevision]= Field(..., discriminator='revisionType')

data = {
    'revisionType': 'GitHub',
    'gitHubLocation': {
        'repository': 'aaa',
        'commitId': 'bbb'
    }
}

# This fails because there is no field named 'revision'
print(RevisionModel.parse_obj(data))

Is there a reasonably clever way to solve this within Pydantic, or do I have to resort to having fields for all the different types and implementing a method to get the revision based on 'revisionType'?



Solution 1:[1]

Your data dictionary does not match the response example. According to your response example it should be like this:

data = {
    "revision": {
        "revisionType": "GitHub",
        "gitHubLocation": {"repository": "aaa", "commitId": "bbb"},
    }
}

That explains the error:

pydantic.error_wrappers.ValidationError: 1 validation error for RevisionModel
revision
  field required (type=value_error.missing)

But now there is a new error:

pydantic.error_wrappers.ValidationError: 3 validation errors for RevisionModel
revision -> revisionType
  unexpected value; permitted: 'String' (type=value_error.const; given=GitHub; permitted=('String',))
revision -> repository
  field required (type=value_error.missing)
revision -> commitId
  field required (type=value_error.missing)

And that is because your GitHubRevision class does not match your response example either. Full example:

from pydantic import BaseModel, Field
from typing import Literal, Union


class GitHubLocation(BaseModel):
    repository: str
    commitId: str


class GitHubRevision(BaseModel):
    revisionType: Literal["GitHub"]
    gitHubLocation: GitHubLocation


class StringRevision(BaseModel):
    revisionType: Literal["String"]


class RevisionModel(BaseModel):
    revision: Union[StringRevision, GitHubRevision] = Field(
        ..., discriminator="revisionType"
    )


data = {
    "revision": {
        "revisionType": "GitHub",
        "gitHubLocation": {"repository": "aaa", "commitId": "bbb"},
    }
}

print(RevisionModel.parse_obj(data))

# Output:
# revision=GitHubRevision(revisionType='GitHub', gitHubLocation=GitHubLocation(repository='aaa', commitId='bbb'))

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1 Hernán Alarcón