Generated models using RootModel to represent Annotated Unions are out of order

veaviticus commented 4 months ago

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 defined

Used 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 Friends and Animals models shouldn't list the other model types in its inheritance to RootModel
the Friends and Animals models should be listed after all the other models (if you simply move them below the Dog model, 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 of Animal, the generated python code works fine, but you obviously lose the discriminator then
If you run with --collapse-root-models then the Animals and Friends models 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 the List[Union[Dog, Cat]] rather than on the Union[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`

veaviticus commented 4 months ago

It seems like this simple patch fixes things for me at least. I think it works fine since we're putting the type hints on the root property of the RootModel already, so we don't need the types in the inheritance declaration as well. The type hints on the root property can do postponed annotation just fine since they aren't part of the object's type declaration

index 0a2810bd..6a42e155 100644
--- a/datamodel_code_generator/model/template/pydantic_v2/RootModel.jinja2
+++ b/datamodel_code_generator/model/template/pydantic_v2/RootModel.jinja2
@@ -10,7 +10,7 @@
 {{ decorator }}
 {% endfor -%}

-class {{ class_name }}({{ base_class }}[{{get_type_hint(fields)}}]):{% if comment is defined %}  # {{ comment }}{% endif %}
+class {{ class_name }}({{ base_class }}):{% if comment is defined %}  # {{ comment }}{% endif %}
 {%- if description %}
     """
     {{ description | indent(4) }}

veaviticus commented 4 months ago

Was just about to file the List discriminator issue I mentioned above, but someone beat me to it: https://github.com/koxudaxi/datamodel-code-generator/issues/1937

koxudaxi / datamodel-code-generator

Generated models using RootModel to represent Annotated Unions are out of order #1921