Generated models using RootModel to represent Annotated Unions are out of order
Describe the bug
When generating models via datamodel-codegen against valid jsonschema created from Pydantic models that rely on postponed annotations and Annotated[Union[]]s, the output pydantic models are "out of order" in the file leading to a python NameError.
To Reproduce
Example python:
from __future__ import annotations
import pydantic
import typing
class Dog(pydantic.BaseModel):
name: typing.Literal['dog'] = pydantic.Field('dog', title='woof')
friends: typing.Optional[typing.List[Animal]] = pydantic.Field(title='Friends', default=[])
class Cat(pydantic.BaseModel):
name: typing.Literal['cat'] = pydantic.Field('cat', title='meow')
friends: typing.Optional[typing.List[Animal]] = pydantic.Field(title='Friends', default=[])
class Bird(pydantic.BaseModel):
name: typing.Literal['bird'] = pydantic.Field('bird', title='bird noise')
friends: typing.Optional[typing.List[Animal]] = pydantic.Field(title='Friends', default=[])
# This is the key bit. This Annotated[Union[]] is a way to use the discriminator on the Union
Animal = typing.Annotated[typing.Union[
Dog, Cat, Bird
], pydantic.Field(title='Any animal', discriminator='name')]
class Zoo(pydantic.BaseModel):
animals: typing.List[Animal] = pydantic.Field(title='A zoo of Animals', default=[])Example schema:
{
"$defs":
{
"Bird":
{
"properties":
{
"name":
{
"const": "bird",
"default": "bird",
"enum":
[
"bird"
],
"title": "bird noise",
"type": "string"
},
"friends":
{
"anyOf":
[
{
"items":
{
"discriminator":
{
"mapping":
{
"bird": "#/$defs/Bird",
"cat": "#/$defs/Cat",
"dog": "#/$defs/Dog"
},
"propertyName": "name"
},
"oneOf":
[
{
"$ref": "#/$defs/Dog"
},
{
"$ref": "#/$defs/Cat"
},
{
"$ref": "#/$defs/Bird"
}
],
"title": "Any animal"
},
"type": "array"
},
{
"type": "null"
}
],
"default":
[],
"title": "Friends"
}
},
"title": "Bird",
"type": "object"
},
"Cat":
{
"properties":
{
"name":
{
"const": "cat",
"default": "cat",
"enum":
[
"cat"
],
"title": "meow",
"type": "string"
},
"friends":
{
"anyOf":
[
{
"items":
{
"discriminator":
{
"mapping":
{
"bird": "#/$defs/Bird",
"cat": "#/$defs/Cat",
"dog": "#/$defs/Dog"
},
"propertyName": "name"
},
"oneOf":
[
{
"$ref": "#/$defs/Dog"
},
{
"$ref": "#/$defs/Cat"
},
{
"$ref": "#/$defs/Bird"
}
],
"title": "Any animal"
},
"type": "array"
},
{
"type": "null"
}
],
"default":
[],
"title": "Friends"
}
},
"title": "Cat",
"type": "object"
},
"Dog":
{
"properties":
{
"name":
{
"const": "dog",
"default": "dog",
"enum":
[
"dog"
],
"title": "woof",
"type": "string"
},
"friends":
{
"anyOf":
[
{
"items":
{
"discriminator":
{
"mapping":
{
"bird": "#/$defs/Bird",
"cat": "#/$defs/Cat",
"dog": "#/$defs/Dog"
},
"propertyName": "name"
},
"oneOf":
[
{
"$ref": "#/$defs/Dog"
},
{
"$ref": "#/$defs/Cat"
},
{
"$ref": "#/$defs/Bird"
}
],
"title": "Any animal"
},
"type": "array"
},
{
"type": "null"
}
],
"default":
[],
"title": "Friends"
}
},
"title": "Dog",
"type": "object"
}
},
"properties":
{
"animals":
{
"default":
[],
"items":
{
"discriminator":
{
"mapping":
{
"bird": "#/$defs/Bird",
"cat": "#/$defs/Cat",
"dog": "#/$defs/Dog"
},
"propertyName": "name"
},
"oneOf":
[
{
"$ref": "#/$defs/Dog"
},
{
"$ref": "#/$defs/Cat"
},
{
"$ref": "#/$defs/Bird"
}
],
"title": "Any animal"
},
"title": "A zoo of Animals",
"type": "array"
}
},
"title": "Zoo",
"type": "object"
}Resulting pydantic from datamodel-codegen
# generated by datamodel-codegen:
# filename: test_pydantic_openapi.json
# timestamp: 2024-04-16T21:54:34+00:00
from __future__ import annotations
from enum import Enum
from typing import List, Literal, Optional, Union
from pydantic import BaseModel, Field, RootModel
class Name(Enum):
bird = 'bird'
class Name1(Enum):
cat = 'cat'
class Name2(Enum):
dog = 'dog'
class Animals(RootModel[Union[Dog, Cat, Bird]]):
root: Union[Dog, Cat, Bird] = Field(..., discriminator='name', title='Any animal')
class Zoo(BaseModel):
animals: Optional[List[Animals]] = Field([], title='A zoo of Animals')
class Friends(RootModel[Union[Dog, Cat, Bird]]):
root: Union[Dog, Cat, Bird] = Field(..., discriminator='name', title='Any animal')
class Bird(BaseModel):
name: Literal['bird'] = Field('bird', title='bird noise')
friends: Optional[List[Friends]] = Field([], title='Friends')
class Cat(BaseModel):
name: Literal['cat'] = Field('cat', title='meow')
friends: Optional[List[Friends]] = Field([], title='Friends')
class Dog(BaseModel):
name: Literal['dog'] = Field('dog', title='woof')
friends: Optional[List[Friends]] = Field([], title='Friends')
Animals.model_rebuild()
Friends.model_rebuild()Attempt to use:
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "test_pydantic_openapi_models.py", line 25, in <module>
class Animals(RootModel[Union[Dog, Cat, Bird]]):
NameError: name 'Dog' is not definedUsed commandline:
$ datamodel-codegen --input test_pydantic_openapi.json --output test_pydantic_openapi_models.py --output-model-type 'pydantic_v2.BaseModel' --field-constraints --target-python-version '3.8'
Expected behavior
In the generated output, one of:
- the
FriendsandAnimalsmodels shouldn't list the other model types in its inheritance toRootModel - the
FriendsandAnimalsmodels should be listed after all the other models (if you simply move them below theDogmodel, everything works fine due to the postponed annotations)
Version:
- OS: Linux / Rocky 8
- Python version: 3.10
- datamodel-code-generator version: 0.25.2
Additional context
- This only seems to happen when the models contain a reference to themselves as part of an Annotated Union. If you directly inline the
Union[Dog,Cat,Bird]into each spot in place ofAnimal, the generated python code works fine, but you obviously lose thediscriminatorthen - If you run with
--collapse-root-modelsthen theAnimalsandFriendsmodels aren't generated at all (are inlined to the other classes) but you end up with a different bug (I'll write up a report for that as well), where the discriminator is placed incorrectly and is attempting to discriminate on theList[Union[Dog, Cat]]rather than on theUnion[Dog,Cat]itself, but in summary:
# generated by datamodel-codegen:
# filename: test_pydantic_openapi.json
# timestamp: 2024-04-16T22:17:08+00:00
from __future__ import annotations
from enum import Enum
from typing import List, Literal, Optional, Union
from pydantic import BaseModel, Field
class Name(Enum):
bird = 'bird'
class Name1(Enum):
cat = 'cat'
class Name2(Enum):
dog = 'dog'
class Zoo(BaseModel):
animals: Optional[List[Union[Dog, Cat, Bird]]] = Field(
[], discriminator='name', title='A zoo of Animals'
)
class Bird(BaseModel):
name: Literal['bird'] = Field('bird', title='bird noise')
friends: Optional[List[Union[Dog, Cat, Bird]]] = Field(
[], discriminator='name', title='Friends'
)
class Cat(BaseModel):
name: Literal['cat'] = Field('cat', title='meow')
friends: Optional[List[Union[Dog, Cat, Bird]]] = Field(
[], discriminator='name', title='Friends'
)
class Dog(BaseModel):
name: Literal['dog'] = Field('dog', title='woof')
friends: Optional[List[Union[Dog, Cat, Bird]]] = Field(
[], discriminator='name', title='Friends'
)so you end up with
TypeError: 'list' is not a valid discriminated union variant; should be a `BaseModel` or `dataclass`
koxudaxi / datamodel-code-generator
Pydantic model and dataclasses.dataclass generator for easy conversion of JSON, OpenAPI, JSON Schema, and YAML data sources.
2.7k
How does funding with Polar work?
1
Pay now to fund the work behind this issue.
2
Get updates on progress being made.
3
Maintainer is rewarded once the issue is completed.
FAQ
Backer
You're funding impactful open source efforts
Contributor
You want to contribute to this effort
Maintainer
You want to get funding like this too