diff --git a/src/datamodel_code_generator/parser/base.py b/src/datamodel_code_generator/parser/base.py index a7ef4412d..73eb59d4f 100644 --- a/src/datamodel_code_generator/parser/base.py +++ b/src/datamodel_code_generator/parser/base.py @@ -110,6 +110,7 @@ def __ge__(self, value: Any, /) -> bool: ... # noqa: D105 ModelNames: TypeAlias = set[ModelName] ModelDeps: TypeAlias = dict[ModelName, set[ModelName]] OrderIndex: TypeAlias = dict[ModelName, int] +DiscriminatorValue: TypeAlias = str | int | bool _BUILTIN_NAMES: frozenset[str] = frozenset(name for name in builtins.__dict__ if not name.startswith("_")) _BUILTIN_NAMES_INTRODUCED_IN: dict[PythonVersion, frozenset[str]] = { @@ -1514,7 +1515,7 @@ def __extract_inherited_enum(cls, models: list[DataModel]) -> None: def _create_discriminator_data_type( self, enum_source: Enum | None, - type_names: list[str], + discriminator_values: list[DiscriminatorValue], discriminator_model: DataModel, imports: Imports, ) -> DataType: @@ -1522,17 +1523,17 @@ def _create_discriminator_data_type( if enum_source: enum_class_name = enum_source.reference.short_name enum_member_literals: list[tuple[str, str]] = [] - for value in type_names: + for value in discriminator_values: member = enum_source.find_member(value) if member and member.field.name: enum_member_literals.append((enum_class_name, member.field.name)) else: # pragma: no cover - enum_member_literals.append((enum_class_name, value)) + enum_member_literals.append((enum_class_name, str(value))) data_type = self.data_type(enum_member_literals=enum_member_literals) if enum_source.module_path != discriminator_model.module_path: # pragma: no cover imports.append(Import.from_full_path(enum_source.name)) else: - data_type = self.data_type(literals=type_names) + data_type = self.data_type(literals=discriminator_values) return data_type def __apply_discriminator_type( # noqa: PLR0912, PLR0914, PLR0915 @@ -1572,12 +1573,12 @@ def __apply_discriminator_type( # noqa: PLR0912, PLR0914, PLR0915 ): # pragma: no cover continue - type_names: list[str] = [] + discriminator_values: list[DiscriminatorValue] = [] def check_paths( model: pydantic_model_v2.BaseModel | Reference, mapping: dict[str, str], - type_names: list[str] = type_names, + discriminator_values: list[DiscriminatorValue] = discriminator_values, ) -> None: """Validate discriminator mapping paths for a model.""" for name, path in mapping.items(): @@ -1589,50 +1590,49 @@ def check_paths( t_disc_2 = "/".join(t_disc.split("/")[1:]) if t_path not in {t_disc, t_disc_2}: # pragma: no branch continue - type_names.append(name) + discriminator_values.append(name) + + def get_discriminator_field_value( + discriminator_field: DataModelFieldBase, + ) -> DiscriminatorValue | None: + const_value = discriminator_field.extras.get("const") + if const_value is not None: + return const_value + + literals = discriminator_field.data_type.literals + if len(literals) == 1: + return literals[0] + + enum_source = discriminator_field.data_type.find_source(Enum) + if enum_source and len(enum_source.fields) == 1: + raw_default = enum_source.fields[0].default + if isinstance(raw_default, str): + return raw_default.strip("'\"") + return raw_default + return None - # First try to get the discriminator value from the const field for discriminator_field in discriminator_model.fields: if field_name not in {discriminator_field.original_name, discriminator_field.name}: continue - if discriminator_field.extras.get("const"): - type_names = [discriminator_field.extras["const"]] + discriminator_value = get_discriminator_field_value(discriminator_field) + if discriminator_value is not None: + discriminator_values = [discriminator_value] break - # If no const value found, try to get it from the mapping - if not type_names: - # Check the main discriminator model path - if mapping: - check_paths(discriminator_model, mapping) # ty: ignore + if not discriminator_values and mapping: + check_paths(discriminator_model, mapping) # ty: ignore - # Check the base_classes if they exist - if len(type_names) == 0: - for base_class in discriminator_model.base_classes: - check_paths(base_class.reference, mapping) # ty: ignore - else: - for discriminator_field in discriminator_model.fields: - if field_name not in {discriminator_field.original_name, discriminator_field.name}: - continue + if len(discriminator_values) == 0: + for base_class in discriminator_model.base_classes: + check_paths(base_class.reference, mapping) # ty: ignore - literals = discriminator_field.data_type.literals - if literals and len(literals) == 1: # pragma: no cover - type_names = [str(v) for v in literals] - break - - enum_source = discriminator_field.data_type.find_source(Enum) - if enum_source and len(enum_source.fields) == 1: - first_field = enum_source.fields[0] - raw_default = first_field.default - if isinstance(raw_default, str): - type_names = [raw_default.strip("'\"")] - else: # pragma: no cover - type_names = [str(raw_default)] - break + if not discriminator_values: + discriminator_values = [discriminator_model.path.split("/")[-1]] - if not type_names: - type_names = [discriminator_model.path.split("/")[-1]] + if not discriminator_values: + discriminator_values = [discriminator_model.path.split("/")[-1]] - if not type_names: # pragma: no cover + if not discriminator_values: # pragma: no cover msg = f"Discriminator type is not found. {data_type.reference.path}" raise RuntimeError(msg) @@ -1666,7 +1666,7 @@ def check_paths( continue literals = discriminator_field.data_type.literals const_value = discriminator_field.extras.get("const") - expected_value = type_names[0] if type_names else None + expected_value = discriminator_values[0] if discriminator_values else None # Check if literals match (existing behavior) literals_match = len(literals) == 1 and literals[0] == expected_value @@ -1701,7 +1701,7 @@ def check_paths( field_data_type.remove_reference() discriminator_field.data_type = self._create_discriminator_data_type( - enum_source, type_names, discriminator_model, imports + enum_source, discriminator_values, discriminator_model, imports ) discriminator_field.data_type.parent = discriminator_field discriminator_field.required = True @@ -1709,7 +1709,7 @@ def check_paths( has_one_literal = True if not has_one_literal: new_data_type = self._create_discriminator_data_type( - enum_from_base, type_names, discriminator_model, imports + enum_from_base, discriminator_values, discriminator_model, imports ) # Handle multiple aliases (Pydantic v2 AliasChoices) single_alias: str | None = None diff --git a/tests/data/expected/main/openapi/discriminator/integer_mapping.py b/tests/data/expected/main/openapi/discriminator/integer_mapping.py new file mode 100644 index 000000000..753634d2f --- /dev/null +++ b/tests/data/expected/main/openapi/discriminator/integer_mapping.py @@ -0,0 +1,30 @@ +# generated by datamodel-codegen: +# filename: discriminator_integer_mapping.yaml +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from enum import IntEnum +from typing import Literal + +from pydantic import BaseModel, Field, RootModel + + +class Kind(IntEnum): + integer_1 = 1 + + +class Foo(BaseModel): + kind: Literal[1] + + +class Kind1(IntEnum): + integer_2 = 2 + + +class Bar(BaseModel): + kind: Literal[2] + + +class Base(RootModel[Foo | Bar]): + root: Foo | Bar = Field(..., discriminator='kind') diff --git a/tests/data/expected/main/openapi/discriminator/integer_no_mapping.py b/tests/data/expected/main/openapi/discriminator/integer_no_mapping.py new file mode 100644 index 000000000..bf9d65d2c --- /dev/null +++ b/tests/data/expected/main/openapi/discriminator/integer_no_mapping.py @@ -0,0 +1,30 @@ +# generated by datamodel-codegen: +# filename: discriminator_integer_no_mapping.yaml +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from enum import IntEnum +from typing import Literal + +from pydantic import BaseModel, Field, RootModel + + +class Kind(IntEnum): + integer_1 = 1 + + +class Foo(BaseModel): + kind: Literal[1] + + +class Kind1(IntEnum): + integer_2 = 2 + + +class Bar(BaseModel): + kind: Literal[2] + + +class Base(RootModel[Foo | Bar]): + root: Foo | Bar = Field(..., discriminator='kind') diff --git a/tests/data/expected/main/openapi/discriminator/integer_no_mapping_literal.py b/tests/data/expected/main/openapi/discriminator/integer_no_mapping_literal.py new file mode 100644 index 000000000..6db573597 --- /dev/null +++ b/tests/data/expected/main/openapi/discriminator/integer_no_mapping_literal.py @@ -0,0 +1,21 @@ +# generated by datamodel-codegen: +# filename: discriminator_integer_no_mapping.yaml +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, Field, RootModel + + +class Foo(BaseModel): + kind: Literal[1] + + +class Bar(BaseModel): + kind: Literal[2] + + +class Base(RootModel[Foo | Bar]): + root: Foo | Bar = Field(..., discriminator='kind') diff --git a/tests/data/expected/main/openapi/discriminator/partial_mapping.py b/tests/data/expected/main/openapi/discriminator/partial_mapping.py new file mode 100644 index 000000000..6c3510b98 --- /dev/null +++ b/tests/data/expected/main/openapi/discriminator/partial_mapping.py @@ -0,0 +1,27 @@ +# generated by datamodel-codegen: +# filename: discriminator_partial_mapping.yaml +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, Field + + +class BaseItem(BaseModel): + itemType: str + + +class FooItem(BaseItem): + fooValue: str | None = None + itemType: Literal['foo'] + + +class BarItem(BaseItem): + barValue: int | None = None + itemType: Literal['BarItem'] + + +class ItemContainer(BaseModel): + item: FooItem | BarItem = Field(..., discriminator='itemType') diff --git a/tests/data/openapi/discriminator_integer_mapping.yaml b/tests/data/openapi/discriminator_integer_mapping.yaml new file mode 100644 index 000000000..da518bffa --- /dev/null +++ b/tests/data/openapi/discriminator_integer_mapping.yaml @@ -0,0 +1,34 @@ +openapi: 3.0.3 +info: + title: Minimal Integer Discriminator + version: 1.0.0 +paths: {} +components: + schemas: + Base: + oneOf: + - $ref: '#/components/schemas/Foo' + - $ref: '#/components/schemas/Bar' + discriminator: + propertyName: kind + mapping: + '1': '#/components/schemas/Foo' + '2': '#/components/schemas/Bar' + Foo: + type: object + properties: + kind: + type: integer + enum: + - 1 + required: + - kind + Bar: + type: object + properties: + kind: + type: integer + enum: + - 2 + required: + - kind diff --git a/tests/data/openapi/discriminator_integer_no_mapping.yaml b/tests/data/openapi/discriminator_integer_no_mapping.yaml new file mode 100644 index 000000000..4aced9a89 --- /dev/null +++ b/tests/data/openapi/discriminator_integer_no_mapping.yaml @@ -0,0 +1,31 @@ +openapi: 3.0.3 +info: + title: Minimal Integer Discriminator + version: 1.0.0 +paths: {} +components: + schemas: + Base: + oneOf: + - $ref: '#/components/schemas/Foo' + - $ref: '#/components/schemas/Bar' + discriminator: + propertyName: kind + Foo: + type: object + properties: + kind: + type: integer + enum: + - 1 + required: + - kind + Bar: + type: object + properties: + kind: + type: integer + enum: + - 2 + required: + - kind diff --git a/tests/data/openapi/discriminator_partial_mapping.yaml b/tests/data/openapi/discriminator_partial_mapping.yaml new file mode 100644 index 000000000..b32a1532d --- /dev/null +++ b/tests/data/openapi/discriminator_partial_mapping.yaml @@ -0,0 +1,48 @@ +openapi: 3.1.0 +info: + title: Test + version: 0.0.0 +paths: + /item: + get: + responses: + '200': + description: Item + content: + application/json: + schema: + $ref: '#/components/schemas/ItemContainer' +components: + schemas: + ItemContainer: + type: object + required: + - item + properties: + item: + $ref: '#/components/schemas/BaseItem' + BaseItem: + type: object + required: + - itemType + properties: + itemType: + type: string + discriminator: + propertyName: itemType + mapping: + foo: '#/components/schemas/FooItem' + FooItem: + allOf: + - $ref: '#/components/schemas/BaseItem' + - type: object + properties: + fooValue: + type: string + BarItem: + allOf: + - $ref: '#/components/schemas/BaseItem' + - type: object + properties: + barValue: + type: integer diff --git a/tests/main/openapi/test_main_openapi.py b/tests/main/openapi/test_main_openapi.py index 26c4db12f..65121180d 100644 --- a/tests/main/openapi/test_main_openapi.py +++ b/tests/main/openapi/test_main_openapi.py @@ -154,6 +154,61 @@ def test_main_openapi_discriminator_enum_duplicate(output_file: Path) -> None: ) +@pytest.mark.skipif( + black.__version__.split(".")[0] == "19", + reason="Installed black doesn't support the old style", +) +def test_main_openapi_discriminator_integer_mapping(output_file: Path) -> None: + """Integer discriminator mapping preserves integer literal values.""" + run_main_and_assert( + input_path=OPEN_API_DATA_PATH / "discriminator_integer_mapping.yaml", + output_path=output_file, + input_file_type="openapi", + assert_func=assert_file_content, + expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "integer_mapping.py", + extra_args=["--target-python-version", "3.10", "--output-model-type", "pydantic_v2.BaseModel"], + ) + + +@pytest.mark.skipif( + black.__version__.split(".")[0] == "19", + reason="Installed black doesn't support the old style", +) +def test_main_openapi_discriminator_integer_no_mapping(output_file: Path) -> None: + """Integer discriminator without mapping preserves integer literal values.""" + run_main_and_assert( + input_path=OPEN_API_DATA_PATH / "discriminator_integer_no_mapping.yaml", + output_path=output_file, + input_file_type="openapi", + assert_func=assert_file_content, + expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "integer_no_mapping.py", + extra_args=["--target-python-version", "3.10", "--output-model-type", "pydantic_v2.BaseModel"], + ) + + +@pytest.mark.skipif( + black.__version__.split(".")[0] == "19", + reason="Installed black doesn't support the old style", +) +def test_main_openapi_discriminator_integer_no_mapping_literal(output_file: Path) -> None: + """Integer discriminator literals remain integers when enums collapse to literals.""" + run_main_and_assert( + input_path=OPEN_API_DATA_PATH / "discriminator_integer_no_mapping.yaml", + output_path=output_file, + input_file_type="openapi", + assert_func=assert_file_content, + expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "integer_no_mapping_literal.py", + extra_args=[ + "--target-python-version", + "3.10", + "--output-model-type", + "pydantic_v2.BaseModel", + "--enum-field-as-literal", + "one", + ], + ) + + @pytest.mark.skipif( black.__version__.split(".")[0] == "19", reason="Installed black doesn't support the old style", @@ -300,6 +355,21 @@ def test_main_openapi_discriminator_short_mapping_names(output_file: Path) -> No ) +def test_main_openapi_discriminator_partial_mapping(output_file: Path) -> None: + """Missing discriminator mappings fall back to the subtype name.""" + run_main_and_assert( + input_path=OPEN_API_DATA_PATH / "discriminator_partial_mapping.yaml", + output_path=output_file, + input_file_type="openapi", + assert_func=assert_file_content, + expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "partial_mapping.py", + extra_args=[ + "--output-model-type", + "pydantic_v2.BaseModel", + ], + ) + + def test_main_openapi_discriminator_no_mapping(output_file: Path) -> None: """Test OpenAPI generation with discriminator without mapping.