Missing Imports When Model Is Changed as a Discriminator

[luca-knaack-webcom]

Describe the bug int Parser.parse function the self.__apply_discriminator_type(model, imports) function gets called for module, models in module_models, before module, models, init, imports, scoped_model_resolver get added to the processed_models array. If certain condition match, this function call has side effects on the already processed models, changing type annotations in the model but not adjusting the imports, thus leading to incorrect models.

To Reproduce

Example schemas schema.json:

{
  "properties": {
    "inner": {
      "discriminator": {
        "mapping": {
          "a": "./type_1.json",
          "A": "./type_1.json"
        },
        "propertyName": "type_"
      },
      "oneOf": [
        {
          "$ref": "./type_1.json"
        }
      ],
      "title": "Inner"
    }
  },
  "required": [
    "inner"
  ],
  "title": "Response",
  "type": "object"
}

type_1.json:

{
  "properties": {
    "type_": {
      "default": "a",
      "enum": ["a", "A"],

      "type": "string",
      "title": "Type"
    }
  },
  "title": "Type1",
  "type": "object"
}

Used commandline:

$ datamodel-codegen --input folder/where/the/files/are --output /output --output-model-type pydantic_v2.BaseModel

Expected behavior I would expect the resulting pydantic files to import everything they use, but for type_1.json I get:

from __future__ import annotations

from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field


class Type(Enum):
    a = 'a'
    A = 'A'


class Type1(BaseModel):
    type_: Literal['a', 'A'] = Field(..., title='Type')

This model imports Optional, which is not used anymore. Optional was deleted here, while handling schema.json. Before the __apply_discriminator_type method call for schema.json, the model for `type_1.json was:

class Type(Enum):
    a = 'a'
    A = 'A'
class Type1(BaseModel):
    type_: Optional[Type] = Field('a', title='Type')

and after:

class Type(Enum):
    a = 'a'
    A = 'A'
class Type1(BaseModel):
    type_: Literal['a', 'A'] = Field(..., title='Type')

Note, that the Type1 model of type_1.json in the processed_models contains the missing import, but also the wrong Optional import.

Version:

• OS: macOS
• Python version: 3.12.3
• datamodel-code-generator version: using local version with commit 5727116

Additional context The missing import can be easily added by posprocessing the processed_models:

for processed_model in processed_models:
    for model in processed_model.models:
        processed_model.imports.append(model.imports)

However, this does not remove the unused model (e.g. Type in this example) or unused imports (e.g. Optional).