self-referencing object with discriminator produces undefined variables for non-pydantic v1 output type

[verngutz]

Describe the bug An OpenAPI schema with a self-referencing object with a discriminated union produces code with undefined variables for all output model types except pydantic v1 BaseModel.

To Reproduce

Example schema:

---
openapi: 3.1.0
components:
  schemas:
    Bar:
      type: object
      properties:
        kind:
          const: Bar
      required:
      - kind
    Foo:
      type: object
      properties:
        kind:
          const: Foo
        children:
          type: array
          items:
            oneOf:
            - "$ref": "#/components/schemas/Foo"
            - "$ref": "#/components/schemas/Bar"
            discriminator:
              propertyName: kind
              mapping:
                Foo: "#/components/schemas/Foo"
                Bar: "#/components/schemas/Bar"
          title: Children
      required:
      - children
      - kind

Used commandline:

$ datamodel-codegen --output-model-type typing.TypedDict

Output:

from __future__ import annotations

from typing import List, Union

from typing_extensions import NotRequired, TypedDict


class Bar(TypedDict):
    kind: NotRequired[str]


Children = Union[Foo, Bar]  # Foo is undefined!


class Foo(TypedDict):
    kind: NotRequired[str]
    children: List[Children]

Similar error happens for other output model types. Only pydantic v1 works because Foo is a variable annotation, rather than in a Union in a type alias.

Expected behavior Reference to Foo properly turned into forward reference

Version:

• OS: Ubuntu 20.04
• Python version: 3.12
• datamodel-code-generator version: 0.25.3

Additional context It works properly if the discriminator is removed from the schema.

[lmeninato]

Would be great to have a fix for this - a simple fix might be if there's a forward reference (https://peps.python.org/pep-0484/#forward-references) just wrap it in quotes.

{
    "$id": "schema",
    "definitions": {
        "Foo": {
            "oneOf": [
                {
                    "type": "string"
                },
                {
                    "type": "array",
                    "items": {
                        "$ref": "#/definitions/Foo"
                    }
                }
            ]
        }
    },
    "title": "State",
    "oneOf": [
        { "$ref": "#/definitions/Foo" }
    ]
}

Currently generates:

# generated by datamodel-codegen:
#   filename:  schema.json
#   timestamp: 2024-02-23T17:00:22+00:00

from __future__ import annotations

from typing import List, Union

from pydantic import Field, RootModel


class Foo(RootModel[Union[str, List[Foo]]]):
    root: Union[str, List[Foo]]


class State(RootModel[Foo]):
    root: Foo = Field(..., title='State')


Foo.model_rebuild()

But simply adding quotes around the Foo in the discriminated union would fix the issue:

# generated by datamodel-codegen:
#   filename:  schema.json
#   timestamp: 2024-02-23T17:00:22+00:00

from __future__ import annotations

from typing import List, Union

from pydantic import Field, RootModel


class Foo(RootModel[Union[str, List["Foo"]]]):
    root: Union[str, List[Foo]]


class State(RootModel[Foo]):
    root: Foo = Field(..., title='State')


Foo.model_rebuild()