Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 41 additions & 0 deletions src/datamodel_code_generator/model/pydantic_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,47 @@ def _get_default_as_pydantic_model(self) -> str | None: # noqa: PLR0911, PLR091
f"lambda :[{data_type_child.alias or data_type_child.reference.source.class_name}."
f"{self._PARSE_METHOD}(v) for v in {self.default!r}]"
)
# If the field references a type alias whose underlying type is a union of BaseModels,
# use TypeAdapter to validate the default value against the full union.
if (
not self.data_type.data_types
and self.data_type.reference
and isinstance(self.data_type.reference.source, DataModel)
and getattr(self.data_type.reference.source, "IS_ALIAS", False)
and isinstance(self.default, (dict, list))
):
alias_source = self.data_type.reference.source
if alias_source.fields and alias_source.fields[0].data_type.data_types:
has_base_model = any(
dt.reference and isinstance(dt.reference.source, BaseModelBase)
for dt in alias_source.fields[0].data_type.data_types
)
if has_base_model:
alias_name = self.data_type.alias or alias_source.class_name
self.__dict__["_type_adapter_import_needed"] = True
return f"lambda :TypeAdapter({alias_name}).validate_python({self.default!r})"

# Inline union with multiple BaseModel candidates: validate against the full union via TypeAdapter
if (
self.data_type.is_union
and isinstance(self.default, dict)
and not any(dt.is_dict for dt in self.data_type.data_types)
):
base_model_count = sum(
1 for dt in self.data_type.data_types if dt.reference and isinstance(dt.reference.source, BaseModelBase)
)
if base_model_count > 1:
type_parts = []
for dt in self.data_type.data_types:
hint = dt.type_hint
if hint and hint != "None":
type_parts.append(hint)
if type_parts:
union_expr = ", ".join(type_parts)
self.__dict__["_type_adapter_import_needed"] = True
self.__dict__["_union_import_for_type_adapter"] = True
return f"lambda :TypeAdapter(Union[{union_expr}]).validate_python({self.default!r})"

for data_type in self.data_type.data_types or (self.data_type,):
# TODO: Check nested data_types
if data_type.is_dict:
Expand Down
8 changes: 7 additions & 1 deletion src/datamodel_code_generator/model/pydantic_v2/base_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

from pydantic import Field, field_validator, model_validator

from datamodel_code_generator.imports import IMPORT_ANY, Import
from datamodel_code_generator.imports import IMPORT_ANY, IMPORT_UNION, Import
from datamodel_code_generator.model.base import ALL_MODEL, UNDEFINED, BaseClassDataType, DataModelFieldBase
from datamodel_code_generator.model.imports import IMPORT_CLASSVAR
from datamodel_code_generator.model.pydantic_base import (
Expand Down Expand Up @@ -200,6 +200,12 @@ def imports(self) -> tuple[Import, ...]:
extra_imports.append(IMPORT_ALIAS_CHOICES)
if self._has_discriminator_in_data_type():
extra_imports.append(IMPORT_FIELD)
if getattr(self, "_type_adapter_import_needed", False):
from datamodel_code_generator.model.pydantic_v2.imports import IMPORT_TYPE_ADAPTER # noqa: PLC0415

extra_imports.append(IMPORT_TYPE_ADAPTER)
if getattr(self, "_union_import_for_type_adapter", False):
extra_imports.append(IMPORT_UNION)
if extra_imports:
return chain_as_tuple(base_imports, tuple(extra_imports))
return base_imports
Expand Down
1 change: 1 addition & 0 deletions src/datamodel_code_generator/model/pydantic_v2/imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
IMPORT_SERIALIZE_AS_ANY = Import.from_full_path("pydantic.SerializeAsAny")
IMPORT_PYDANTIC_DATACLASS = Import.from_full_path("pydantic.dataclasses.dataclass")
IMPORT_ROOT_MODEL = Import.from_full_path("pydantic.RootModel")
IMPORT_TYPE_ADAPTER = Import.from_full_path("pydantic.TypeAdapter")
IMPORT_FIELD_VALIDATOR = Import.from_full_path("pydantic.field_validator")
IMPORT_VALIDATION_INFO = Import.from_full_path("pydantic.ValidationInfo")
IMPORT_VALIDATOR_FUNCTION_WRAP_HANDLER = Import.from_full_path("pydantic.ValidatorFunctionWrapHandler")
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
# generated by datamodel-codegen:
# filename: multi_model_union_default_object.json
# timestamp: 2019-07-26T00:00:00+00:00

from __future__ import annotations

from typing import Annotated, Literal, Union

from pydantic import BaseModel, Field, TypeAdapter


class A(BaseModel):
type: Literal['a']


class B(BaseModel):
type: Literal['b']


class Model(BaseModel):
x: Annotated[
A | B | None,
Field(
default_factory=lambda: TypeAdapter(Union[A, B]).validate_python(
{'type': 'b'}
)
),
]
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
# generated by datamodel-codegen:
# filename: multi_model_union_default_object.json
# timestamp: 2019-07-26T00:00:00+00:00

from __future__ import annotations

from typing import Annotated, Literal, Union

from pydantic import BaseModel, Field, TypeAdapter


class A(BaseModel):
type: Literal['a']


class B(BaseModel):
type: Literal['b']


class Model(BaseModel):
x: Annotated[
A | B | None,
Field(
default_factory=lambda: TypeAdapter(Union[A, B]).validate_python(
{'type': 'b'}
)
),
]
27 changes: 27 additions & 0 deletions tests/data/expected/main/jsonschema/ref_union_default_object.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
# generated by datamodel-codegen:
# filename: ref_union_default_object.json
# timestamp: 2019-07-26T00:00:00+00:00

from __future__ import annotations

from typing import Annotated, Literal

from pydantic import BaseModel, Field, TypeAdapter


class A(BaseModel):
type: Literal['a']


class B(BaseModel):
type: Literal['b']


type U = A | B


class Model(BaseModel):
x: Annotated[
U | None,
Field(default_factory=lambda: TypeAdapter(U).validate_python({'type': 'b'})),
]
35 changes: 35 additions & 0 deletions tests/data/jsonschema/multi_model_union_default_object.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"$defs": {
"A": {
"type": "object",
"properties": {
"type": {
"const": "a",
"type": "string"
}
},
"required": ["type"]
},
"B": {
"type": "object",
"properties": {
"type": {
"const": "b",
"type": "string"
}
},
"required": ["type"]
}
},
"properties": {
"x": {
"anyOf": [
{ "$ref": "#/$defs/A" },
{ "$ref": "#/$defs/B" }
],
"default": { "type": "b" }
}
}
}
38 changes: 38 additions & 0 deletions tests/data/jsonschema/ref_union_default_object.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"$defs": {
"A": {
"type": "object",
"properties": {
"type": {
"const": "a",
"type": "string"
}
},
"required": ["type"]
},
"B": {
"type": "object",
"properties": {
"type": {
"const": "b",
"type": "string"
}
},
"required": ["type"]
},
"U": {
"anyOf": [
{ "$ref": "#/$defs/A" },
{ "$ref": "#/$defs/B" }
]
}
},
"properties": {
"x": {
"$ref": "#/$defs/U",
"default": { "type": "b" }
}
}
}
68 changes: 68 additions & 0 deletions tests/main/jsonschema/test_main_jsonschema.py
Original file line number Diff line number Diff line change
Expand Up @@ -5591,6 +5591,74 @@ def test_main_jsonschema_enum_literal_type_alias_default(output_file: Path) -> N
)


@pytest.mark.skipif(
int(black.__version__.split(".")[0]) < 23,
reason="Installed black doesn't support the new 'type' statement",
)
def test_main_jsonschema_ref_union_default_object(output_file: Path) -> None:
"""TypeAlias ref to union of BaseModels uses TypeAdapter for dict default (#3034)."""
run_main_and_assert(
input_path=JSON_SCHEMA_DATA_PATH / "ref_union_default_object.json",
output_path=output_file,
input_file_type=None,
assert_func=assert_file_content,
expected_file="ref_union_default_object.py",
extra_args=[
"--use-type-alias",
"--use-annotated",
"--target-python-version",
"3.12",
"--output-model-type",
"pydantic_v2.BaseModel",
],
)


@pytest.mark.skipif(
int(black.__version__.split(".")[0]) < 23,
reason="Installed black doesn't support the new 'type' statement",
)
def test_main_jsonschema_pydantic_v2_multi_model_union_default_object_type_alias(output_file: Path) -> None:
"""Inline union of multiple BaseModels uses TypeAdapter with --use-type-alias (#3035)."""
run_main_and_assert(
input_path=JSON_SCHEMA_DATA_PATH / "multi_model_union_default_object.json",
output_path=output_file,
input_file_type=None,
assert_func=assert_file_content,
expected_file="pydantic_v2_multi_model_union_default_object_type_alias.py",
extra_args=[
"--use-type-alias",
"--use-annotated",
"--target-python-version",
"3.12",
"--output-model-type",
"pydantic_v2.BaseModel",
],
)


@pytest.mark.skipif(
int(black.__version__.split(".")[0]) < 23,
reason="Installed black doesn't support the new 'type' statement",
)
def test_main_jsonschema_pydantic_v2_multi_model_union_default_object(output_file: Path) -> None:
"""Inline union of multiple BaseModels uses TypeAdapter without --use-type-alias (#3035)."""
run_main_and_assert(
input_path=JSON_SCHEMA_DATA_PATH / "multi_model_union_default_object.json",
output_path=output_file,
input_file_type=None,
assert_func=assert_file_content,
expected_file="pydantic_v2_multi_model_union_default_object.py",
extra_args=[
"--use-annotated",
"--target-python-version",
"3.12",
"--output-model-type",
"pydantic_v2.BaseModel",
],
)


@pytest.mark.cli_doc(
options=["--type-mappings"],
option_description="""Override default type mappings for schema formats.
Expand Down
Loading