Refine discriminator enum mapping

Author: koxudaxi

src/datamodel_code_generator/parser/base.py

@@ -1564,7 +1564,7 @@ def _create_discriminator_data_type(
         if enum_source:
             if self.use_enum_values_in_discriminator:
                 enum_class_name = enum_source.reference.short_name
-                enum_member_literals: list[tuple[str, DiscriminatorValue]] = []
+                enum_member_literals: list[tuple[str, str]] = []
                 for value in discriminator_values:
                     member = enum_source.find_member(value)
                     if member and member.field.name:
@@ -1575,24 +1575,28 @@ def _create_discriminator_data_type(
                 if enum_source.module_path != discriminator_model.module_path:  # pragma: no cover
                     imports.append(Import.from_full_path(enum_source.name))
             else:
-                # According to OpenAPI specification, mapping discriminators are always string values.
-                # However, if the mapped object is an enum, we want to use the real enum value instead of
-                # the string value.
-                # See: https://swagger.io/specification/#options-for-mapping-values-to-schemas
-                # Fix: https://github.com/koxudaxi/datamodel-code-generator/issues/3073
-                for i, value in enumerate(discriminator_values):
-                    if member := enum_source.find_member(value):
-                        match member.field.default:
-                            case str():
-                                discriminator_values[i] = member.field.default.strip("'\"")
-                            case _ if isinstance(member.field.default, DiscriminatorValue):
-                                discriminator_values[i] = member.field.default
-
-                data_type = self.data_type(literals=discriminator_values)
+                data_type = self.data_type(
+                    literals=[
+                        Parser._get_enum_discriminator_literal(enum_source, value) for value in discriminator_values
+                    ]
+                )
         else:
             data_type = self.data_type(literals=discriminator_values)
         return data_type
 
+    @staticmethod
+    def _get_enum_discriminator_literal(enum_source: Enum, value: DiscriminatorValue) -> DiscriminatorValue:
+        member = enum_source.find_member(value)
+        if not member:
+            return value
+
+        default = member.field.default
+        if isinstance(default, str):
+            return default.strip("'\"")
+        if isinstance(default, int | bool):
+            return default
+        return value
+
     def __apply_discriminator_type(  # noqa: PLR0912, PLR0914, PLR0915
         self,
         models: list[DataModel],
@@ -1706,28 +1710,29 @@ def get_discriminator_field_value(
                         msg = f"Discriminator type is not found. {data_type.reference.path}"
                         raise RuntimeError(msg)
 
-                    enum_from_base: Enum | None = None
-                    for base_class in discriminator_model.base_classes:
-                        if not base_class.reference or not base_class.reference.source:  # pragma: no cover
-                            continue
-                        base_model = base_class.reference.source
-                        if not isinstance(  # pragma: no cover
-                            base_model,
-                            (
-                                pydantic_model_v2.BaseModel,
-                                dataclass_model.DataClass,
-                                msgspec_model.Struct,
-                            ),
-                        ):
-                            continue
-                        for base_field in base_model.fields:  # pragma: no branch
-                            if field_name not in {base_field.original_name, base_field.name}:  # pragma: no cover
+                    def get_enum_from_base(
+                        discriminator_model: DataModel = discriminator_model,
+                        field_name: str = field_name,
+                    ) -> Enum | None:
+                        for base_class in discriminator_model.base_classes:
+                            if not base_class.reference or not base_class.reference.source:  # pragma: no cover
                                 continue
-                            enum_from_base = base_field.data_type.find_source(Enum)
-                            if enum_from_base:  # pragma: no branch
-                                break
-                        if enum_from_base:  # pragma: no branch
-                            break
+                            base_model = base_class.reference.source
+                            if not isinstance(  # pragma: no cover
+                                base_model,
+                                (
+                                    pydantic_model_v2.BaseModel,
+                                    dataclass_model.DataClass,
+                                    msgspec_model.Struct,
+                                ),
+                            ):
+                                continue
+                            for base_field in base_model.fields:  # pragma: no branch
+                                if field_name not in {base_field.original_name, base_field.name}:  # pragma: no cover
+                                    continue
+                                if enum_from_base := base_field.data_type.find_source(Enum):  # pragma: no branch
+                                    return enum_from_base
+                        return None
 
                     has_one_literal = False
                     for discriminator_field in discriminator_model.fields:
@@ -1759,7 +1764,9 @@ def get_discriminator_field_value(
                             discriminator_field.extras["is_classvar"] = True
                             break
 
-                        enum_source = discriminator_field.data_type.find_source(Enum) or enum_from_base
+                        enum_source = discriminator_field.data_type.find_source(Enum)
+                        if self.use_enum_values_in_discriminator:
+                            enum_source = enum_source or get_enum_from_base()
 
                         for field_data_type in discriminator_field.data_type.all_data_types:
                             if field_data_type.reference:  # pragma: no cover
@@ -1774,7 +1781,7 @@ def get_discriminator_field_value(
                         has_one_literal = True
                     if not has_one_literal:
                         new_data_type = self._create_discriminator_data_type(
-                            enum_from_base, discriminator_values, discriminator_model, imports
+                            get_enum_from_base(), discriminator_values, discriminator_model, imports
                         )
                         # Handle multiple aliases (Pydantic v2 AliasChoices)
                         single_alias: str | None = None

tests/data/expected/main/openapi/discriminator/float_mapping.py

@@ -0,0 +1,32 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_float_mapping.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Literal
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class Kind(Enum):
+    number_1_5 = 1.5
+    number_2_5 = 2.5
+
+
+class Foo(BaseModel):
+    kind: Literal['1.5']
+
+
+class Kind1(Enum):
+    number_3_5 = 3.5
+    number_4_5 = 4.5
+
+
+class Bar(BaseModel):
+    kind: Literal['3.5']
+
+
+class Base(RootModel[Foo | Bar]):
+    root: Foo | Bar = Field(..., discriminator='kind')

tests/data/expected/main/openapi/discriminator/integer_mapping_use_enum.py

@@ -0,0 +1,31 @@
+# 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[Kind.integer_1]
+
+
+class Kind1(IntEnum):
+    integer_2 = 2
+    integer_3 = 3
+
+
+class Bar(BaseModel):
+    kind: Literal[Kind1.integer_2, Kind1.integer_3]
+
+
+class Base(RootModel[Foo | Bar]):
+    root: Foo | Bar = Field(..., discriminator='kind')

tests/data/openapi/discriminator_float_mapping.yaml

@@ -0,0 +1,36 @@
+openapi: 3.0.3
+info:
+  title: Float Discriminator Mapping
+  version: 1.0.0
+paths: {}
+components:
+  schemas:
+    Base:
+      oneOf:
+        - $ref: '#/components/schemas/Foo'
+        - $ref: '#/components/schemas/Bar'
+      discriminator:
+        propertyName: kind
+        mapping:
+          '1.5': '#/components/schemas/Foo'
+          '3.5': '#/components/schemas/Bar'
+    Foo:
+      type: object
+      properties:
+        kind:
+          type: number
+          enum:
+            - 1.5
+            - 2.5
+      required:
+        - kind
+    Bar:
+      type: object
+      properties:
+        kind:
+          type: number
+          enum:
+            - 3.5
+            - 4.5
+      required:
+        - kind

tests/main/openapi/test_main_openapi.py

@@ -178,6 +178,44 @@ def test_main_openapi_discriminator_integer_mapping(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_use_enum(output_file: Path) -> None:
+    """Integer discriminator mapping preserves enum member 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_use_enum.py",
+        extra_args=[
+            "--target-python-version",
+            "3.10",
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+            "--use-enum-values-in-discriminator",
+        ],
+    )
+
+
+@pytest.mark.skipif(
+    black.__version__.split(".")[0] == "19",
+    reason="Installed black doesn't support the old style",
+)
+def test_main_openapi_discriminator_float_mapping(output_file: Path) -> None:
+    """Unsupported discriminator enum values keep their mapping string values."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_float_mapping.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "float_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",