JSON schema generation of Discriminated Unions doesn't work with lists

[adrienyhuel]

Checks

• I added a descriptive title to this issue
• I have searched (google, github) for similar issues and couldn't find anything
• I have read and followed the docs and still think this is a bug

Bug

Output of python -c "import pydantic.utils; print(pydantic.utils.version_info())":

pydantic version: 1.9.0
            pydantic compiled: True
                 install path: C:\Users\Adrien\Documents\GitHub\mft\backend\venv\Lib\site-packages\pydantic
               python version: 3.9.5 (tags/v3.9.5:0a7dcbd, May  3 2021, 17:27:52) [MSC v.1928 64 bit (AMD64)]
                     platform: Windows-10-10.0.19044-SP0
     optional deps. installed: ['dotenv', 'email-validator', 'typing-extensions']

I tested the ne feature "discriminated unions", using the sample at https://pydantic-docs.helpmanual.io/usage/types/#discriminated-unions-aka-tagged-unions, and it works great.

But when I use it in arrays, it doesn't work. The resulting json schema doesn't contains any discriminator.

See here, I modified the sample ( pet: Pet to pet: List[Pet] ) :

from typing import Literal, Union, List
from typing_extensions import Annotated
from pydantic import BaseModel, Field

class BlackCat(BaseModel):
    pet_type: Literal['cat']
    color: Literal['black']
    black_name: str

class WhiteCat(BaseModel):
    pet_type: Literal['cat']
    color: Literal['white']
    white_name: str

Cat = Annotated[Union[BlackCat, WhiteCat], Field(discriminator='color')]

class Dog(BaseModel):
    pet_type: Literal['dog']
    name: str

Pet = Annotated[Union[Cat, Dog], Field(discriminator='pet_type')]

class Model(BaseModel):
    pet: List[Pet]
    n: int

print(Model.schema_json())

...

JSON schema gives :

"pet":{
         "title":"Pet",
         "type":"array",
         "items":{
            "anyOf":[
               {
                  "anyOf":[
                     {
                        "$ref":"#/definitions/BlackCat"
                     },
                     {
                        "$ref":"#/definitions/WhiteCat"
                     }
                  ]
               },
               {
                  "$ref":"#/definitions/Dog"
               }
            ]
         }
      }

[janluke]

Maybe the title should be more specific. The parsing works correctly AFAICS. The JSON schema generation doesn't.

[nayef-livio-derwiche]

Hi all, this workaround seems to work for me, until using the Annotation directly is fixed for lists:

from typing import Literal, Annotated, Union
from pydantic import BaseModel, Field


class Dog(BaseModel):
    animal_type: Literal['dog']
    name: str
    claw: str

class Bird(BaseModel):
    animal_type: Literal['bird']
    name: str
    wing_size: str


class Animal(BaseModel):
    __root__: Annotated[
        Union[
            Dog,
            Bird
        ],
        Field(discriminator='animal_type')
    ]


# Animal = Annotated[
#     Union[
#         Dog,
#         Bird
#     ],
#     Field(discriminator='animal_type')
# ]

class Zoo(BaseModel):
    animals: list[Animal]


print(Zoo.schema_json())

The schema I get:

{
  "title": "Zoo",
  "type": "object",
  "properties": {
    "animals": {
      "title": "Animals",
      "type": "array",
      "items": {
        "$ref": "#/definitions/Animal"
      }
    }
  },
  "required": [
    "animals"
  ],
  "definitions": {
    "Dog": {
      "title": "Dog",
      "type": "object",
      "properties": {
        "animal_type": {
          "title": "Animal Type",
          "enum": [
            "dog"
          ],
          "type": "string"
        },
        "name": {
          "title": "Name",
          "type": "string"
        },
        "claw": {
          "title": "Claw",
          "type": "string"
        }
      },
      "required": [
        "animal_type",
        "name",
        "claw"
      ]
    },
    "Bird": {
      "title": "Bird",
      "type": "object",
      "properties": {
        "animal_type": {
          "title": "Animal Type",
          "enum": [
            "bird"
          ],
          "type": "string"
        },
        "name": {
          "title": "Name",
          "type": "string"
        },
        "wing_size": {
          "title": "Wing Size",
          "type": "string"
        }
      },
      "required": [
        "animal_type",
        "name",
        "wing_size"
      ]
    },
    "Animal": {
      "title": "Animal",
      "discriminator": {
        "propertyName": "animal_type",
        "mapping": {
          "dog": "#/definitions/Dog",
          "bird": "#/definitions/Bird"
        }
      },
      "anyOf": [
        {
          "$ref": "#/definitions/Dog"
        },
        {
          "$ref": "#/definitions/Bird"
        }
      ]
    }
  }
}

[adrienyhuel]

@nayef-livio-derwiche thank, it worked for me !

I just had to add this to the Animal class, to be able to use dot notation for getting attributes:

    def __getattr__(self, item):  # if you want to use '.'
        return self.__root__.__getattribute__(item)

[nayef-livio-derwiche]

@adrienyhuel

Sure, something like that might help if you want to use the root class transparently:

def use_root_class_decorator(cls):
    def new_method_proxy(func):
        def inner(self, *args):
            return func(self.__root__, *args)
        return inner
    
    cls.__getattr__ = new_method_proxy(getattr)
    cls.__setattr__ = new_method_proxy(setattr)
    cls.__delattr__ = new_method_proxy(delattr)
    cls.__getattr__ = new_method_proxy(getattr)
    cls.__bytes__ = new_method_proxy(bytes)
    cls.__str__ = new_method_proxy(str)
    cls.__bool__ = new_method_proxy(bool)
    cls.__dir__ = new_method_proxy(dir)
    cls.__hash__ = new_method_proxy(hash)
    # not sure of the consequences of using / not using the following line...
    # cls.__class__ = property(new_method_proxy(operator.attrgetter("__class__")))
    cls.__eq__ = new_method_proxy(operator.eq)
    cls.__lt__ = new_method_proxy(operator.lt)
    cls.__gt__ = new_method_proxy(operator.gt)
    cls.__ne__ = new_method_proxy(operator.ne)
    cls.__hash__ = new_method_proxy(hash)
    cls.__getitem__ = new_method_proxy(operator.getitem)
    cls.__setitem__ = new_method_proxy(operator.setitem)
    cls.__delitem__ = new_method_proxy(operator.delitem)
    cls.__iter__ = new_method_proxy(iter)
    cls.__len__ = new_method_proxy(len)
    cls.__contains__ = new_method_proxy(operator.contains)
    
    return cls

@use_root_class_decorator
class Animal(BaseModel):
    __root__: Annotated[
        Union[
            Dog,
            Bird
        ],
        Field(discriminator='animal_type')
    ]

Or if you can, you may use this lib: https://github.com/GrahamDumpleton/wrapt

[samuelcolvin]

Thanks for the patience, I've submitted a fix for this, see #4071.

@adrienyhuel, please could you take a look and test that it fixes your problem. I hope to get v1.9.1 out with this included in the next few days.

[adrienyhuel]

@samuelcolvin
I installed with pip from "git+https://github.com/samuelcolvin/pydantic.git@refs/pull/4071/merge" and it works great.

Thanks for the fix, waiting for 1.9.1 release :)

[samuelcolvin]

Thanks so much for confirming. Should be very soon now.