koxudaxi / datamodel-code-generator

Pydantic model and dataclasses.dataclass generator for easy conversion of JSON, OpenAPI, JSON Schema, and YAML data sources.
https://koxudaxi.github.io/datamodel-code-generator/
MIT License
2.76k stars 302 forks source link

Discriminated union with a List: annotation is on the list, not the union #1937

Open colmanhumphrey opened 6 months ago

colmanhumphrey commented 6 months ago

Describe the bug

Generation of list/array of discriminated union annotates the list, not the union. Same problem with = Field(..., discriminator) too if you don't use annotated.

To Reproduce

Example schema:

{
  "$schema": "http://json-schema.org/draft-04/schema#",
  "title": "SomeTest",
  "type": "object",
  "required": [
    "some_array"
  ],
  "properties": {
    "some_array": {
      "type": "array",
      "items": {
        "discriminator": {
          "propertyName": "discrim",
          "mapping": {
            "Type1": "#/definitions/FeatureType1",
            "Type2": "#/definitions/FeatureType2"
          }
        },
        "oneOf": [
          {
          "$ref": "#/definitions/FeatureType1"
        },
          {
          "$ref": "#/definitions/FeatureType2"
        }
        ]
      }
    }
  },
  "definitions": {
    "FeatureType1": {
      "type": "object",
      "properties": {
        "discrim": {
          "type": "string"
        },
        "id": {
          "type": "string"
        }
      }
    },
    "FeatureType2": {
      "type": "object",
      "properties": {
        "discrim": {
          "type": "string"
        },
        "id": {
          "type": "string"
        },
        "something_else": {
          "type": "string"
        }
      }
    }
  }
}

Used commandline:

$ datamodel-codegen --disable-timestamp --use-annotated --collapse-root-models --target-python-version '3.11' --input ../schema.json --input-file-type jsonschema --output whatever.py

This gives:

# generated by datamodel-codegen:
#   filename:  schema.json

from __future__ import annotations

from typing import Annotated, List, Literal, Optional, Union

from pydantic import BaseModel, Field

class FeatureType1(BaseModel):
    discrim: Literal['Type1']
    id: Optional[str] = None

class FeatureType2(BaseModel):
    discrim: Literal['Type2']
    id: Optional[str] = None
    something_else: Optional[str] = None

class SomeTest(BaseModel):
    some_array: Annotated[
        List[Union[FeatureType1, FeatureType2]], Field(discriminator='discrim')
    ]

This doesn't work:

TypeError: `discriminator` can only be used with `Union` type with more than one variant`

Expected behavior

The problem is that the Annotation should be on the Union, not on the list, so the SomeTest class should be:

class SomeTest(BaseModel):
    some_array: List[
        Annotated[Union[FeatureType1, FeatureType2], Field(discriminator='discrim')]
    ]

And now it works. Note that we get more or less the same problem if you remove --use-annotated

Version:

Additional context

Thanks!

quizmoon commented 1 month ago

That behavior was ignored by Pydantic until 2.8.2, but starting 2.9.0 it end up with the error on generated model:

TypeError: 'list' is not a valid discriminated union variant; should be aBaseModelordataclass``

Maybe there is workaround to bypass it?

pcavalar commented 3 weeks ago

This is blocking us from upgrading pydantic, does anyone have any workaround?

quizmoon commented 3 weeks ago

This is blocking us from upgrading pydantic, does anyone have any workaround?

the same issue :/