Discriminated Unions with callable Discriminator fail to create a class dynamically with Optional annotation

[chiselko6]

Initial Checks

• I confirm that I'm using Pydantic V2

Description

In our code we need to convert existing pydantic BaseModels into new classes with optional fields dynamically. We utilize pydantic.create_model for that wrapping field annotations with Optional. This works for fine for fields being discriminated unions defined as Field(discriminator="<property>"), but fails for ones defined as Discriminator(callable) with corresponding Tag annotation.

In the provided example I built new-style and old-style discriminated unions side-by-side for comparison. The error the new-style discriminated union raises is the following, as if NoneType messes up with the inner union:

Traceback (most recent call last):
  File "/Users/chiselko6/dev/pydantic_discriminator/test.py", line 77, in <module>
    create_model(
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/main.py", line 1550, in create_model
    return meta(
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_model_construction.py", line 202, in __new__
    complete_model_class(
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_model_construction.py", line 539, in complete_model_class
    schema = cls.__get_pydantic_core_schema__(cls, handler)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/main.py", line 626, in __get_pydantic_core_schema__
    return handler(source)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_schema_generation_shared.py", line 82, in __call__
    schema = self._handler(source_type)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_generate_schema.py", line 502, in generate_schema
    schema = self._generate_schema_inner(obj)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_generate_schema.py", line 753, in _generate_schema_inner
    return self._model_schema(obj)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_generate_schema.py", line 580, in _model_schema
    {k: self._generate_md_field_schema(k, v, decorators) for k, v in fields.items()},
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_generate_schema.py", line 580, in <dictcomp>
    {k: self._generate_md_field_schema(k, v, decorators) for k, v in fields.items()},
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_generate_schema.py", line 916, in _generate_md_field_schema
    common_field = self._common_field_schema(name, field_info, decorators)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_generate_schema.py", line 1081, in _common_field_schema
    schema = self._apply_annotations(
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_generate_schema.py", line 1820, in _apply_annotations
    schema = get_inner_schema(source_type)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_schema_generation_shared.py", line 82, in __call__
    schema = self._handler(source_type)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/_internal/_generate_schema.py", line 1902, in new_handler
    schema = metadata_get_schema(source, get_inner_schema)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/types.py", line 2827, in __get_pydantic_core_schema__
    return self._convert_schema(original_schema)
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/types.py", line 2848, in _convert_schema
    raise PydanticUserError(
pydantic.errors.PydanticUserError: `Tag` not provided for choice {'type': 'nullable', 'schema': {'type': 'tagged-union', 'choices': {'int': {'type': 'definition-ref', 'schema_ref': '__main__.ValueInt:140660602191040', 'metadata': {'pydantic.internal.tagged_union_tag': 'int'}}, 'str': {'type': 'definition-ref', 'schema_ref': '__main__.ValueStr:140660580617136', 'metadata': {'pydantic.internal.tagged_union_tag': 'str'}}}, 'discriminator': <function model_x_discriminator at 0x7fee190baf70>}} used with `Discriminator`

For further information visit https://errors.pydantic.dev/2.7/u/callable-discriminator-no-tag

Example Code

from typing import Any, Literal, Optional, Union

from pydantic import BaseModel, Discriminator, Field, Tag, create_model
from typing_extensions import Annotated

# Definitions

def model_x_discriminator(v: Any) -> str:
    if isinstance(v, int):
        return "int"
    if isinstance(v, str):
        return "str"
    raise ValueError()


class ValueInt(BaseModel):
    type: Literal["int"] = "int"
    x: int


class ValueStr(BaseModel):
    type: Literal["str"] = "str"
    x: str


# old-style discriminator by class property

OldDiscriminatedType = Annotated[
    Union[ValueInt, ValueStr],
    Field(discriminator="type"),
]


class OldDiscriminatedModel(BaseModel):
    value: OldDiscriminatedType


class OldDiscriminatedModelOptional(BaseModel):
    value: Optional[OldDiscriminatedType]


old_value_field_info = OldDiscriminatedModel.model_fields["value"]
print(
    create_model(
        "OldDiscriminatedModelOptionalTest",
        __base__=OldDiscriminatedModel,
        **{
            "value": (
                Optional[OldDiscriminatedType],
                old_value_field_info,
            )
        },
    )
)
# outputs a new class


# new-style Discriminator and Tag

NewDiscriminatedType = Annotated[
    Union[
        Annotated[ValueInt, Tag("int")],
        Annotated[ValueStr, Tag("str")],
    ],
    Discriminator(model_x_discriminator),
]


class NewDiscriminatedModel(BaseModel):
    value: NewDiscriminatedType


class NewDiscriminatedModelOptional(BaseModel):
    value: Optional[NewDiscriminatedType]


new_value_field_info = NewDiscriminatedModel.model_fields["value"]
print(
    create_model(
        "NewDiscriminatedModelOptionalTest",
        __base__=NewDiscriminatedModel,
        **{
            "value": (
                Optional[NewDiscriminatedType],
                new_value_field_info,
            )
        },
    )
)
# raises an error

Python, Pydantic & OS Version

$ python3 -c "import pydantic.version; print(pydantic.version.version_info())"
             pydantic version: 2.7.1
        pydantic-core version: 2.18.2
          pydantic-core build: profile=release pgo=false
                 install path: /Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic
               python version: 3.9.0 (v3.9.0:9cf6752276, Oct  5 2020, 11:29:23)  [Clang 6.0 (clang-600.0.57)]
                     platform: macOS-10.16-x86_64-i386-64bit
             related packages: typing_extensions-4.11.0
                       commit: unknown

[sydney-runkle]

@chiselko6,

Thanks for the detailed description and reproducible code snippet. Definitely looks like a bug. Can be a bit difficult to track down with our complicated discriminator application logic, but I'll look into a fix for this in 2.8!

[sydney-runkle]

If anyone wants to take a stab in the meantime, be my guest :).