Map OpenAPI Discriminator mapping's values to their real value in enums

src/datamodel_code_generator/parser/base.py

@@ -1562,21 +1562,41 @@ def _create_discriminator_data_type(
     ) -> DataType:
         """Create a data type for discriminator field, using enum literals if available."""
         if enum_source:
-            enum_class_name = enum_source.reference.short_name
-            enum_member_literals: list[tuple[str, str]] = []
-            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, 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))
+            if self.use_enum_values_in_discriminator:
+                enum_class_name = enum_source.reference.short_name
+                enum_member_literals: list[tuple[str, str]] = []
+                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, 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=[
+                        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],
@@ -1690,8 +1710,10 @@ 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
-                    if self.use_enum_values_in_discriminator:
+                    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
@@ -1708,11 +1730,9 @@ def get_discriminator_field_value(
                             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
-                                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
+                                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:
@@ -1744,11 +1764,9 @@ def get_discriminator_field_value(
                             discriminator_field.extras["is_classvar"] = True
                             break
 
-                        enum_source: Enum | None = None
+                        enum_source = discriminator_field.data_type.find_source(Enum)
                         if self.use_enum_values_in_discriminator:
-                            enum_source = (  # pragma: no cover
-                                discriminator_field.data_type.find_source(Enum) or enum_from_base
-                            )
+                            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
@@ -1763,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.py

@@ -20,10 +20,11 @@ class Foo(BaseModel):
 
 class Kind1(IntEnum):
     integer_2 = 2
+    integer_3 = 3
 
 
 class Bar(BaseModel):
-    kind: Literal[2]
+    kind: Literal[2, 3]
 
 
 class Base(RootModel[Foo | Bar]):

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/data/openapi/discriminator_integer_mapping.yaml

@@ -14,6 +14,7 @@ components:
         mapping:
           '1': '#/components/schemas/Foo'
           '2': '#/components/schemas/Bar'
+          '3': '#/components/schemas/Bar'
     Foo:
       type: object
       properties:
@@ -30,5 +31,6 @@ components:
           type: integer
           enum:
             - 2
+            - 3
       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",