Fix discriminator to use enum value instead of model name for single-value enums (#2652)

* Fix discriminator handling for single-value enums and add related tests

* Fix type name handling for single-value literals in discriminator logic

Author: koxudaxi

src/datamodel_code_generator/parser/base.py

@@ -1086,7 +1086,27 @@ def check_paths(
                                 for base_class in discriminator_model.base_classes:
                                     check_paths(base_class.reference, mapping)  # pyright: ignore[reportArgumentType]
                         else:
-                            type_names = [discriminator_model.path.split("/")[-1]]
+                            for discriminator_field in discriminator_model.fields:
+                                if field_name not in {discriminator_field.original_name, discriminator_field.name}:
+                                    continue
+
+                                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 type_names:
+                                type_names = [discriminator_model.path.split("/")[-1]]
 
                     if not type_names:  # pragma: no cover
                         msg = f"Discriminator type is not found. {data_type.reference.path}"

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

@@ -0,0 +1,25 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_enum_single_value.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class ToolType(Enum):
+    function = 'function'
+
+
+class ToolBase(BaseModel):
+    type: ToolType
+
+
+class FunctionToolCall(ToolBase):
+    id: str
+
+
+class ToolCall(RootModel[FunctionToolCall]):
+    root: FunctionToolCall = Field(..., discriminator='type')

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

@@ -0,0 +1,27 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_enum_single_value_anyof.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Literal, Union
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class ToolType(Enum):
+    function = 'function'
+
+
+class FunctionToolCall(BaseModel):
+    id: str
+    type: Literal['function']
+
+
+class CustomToolCall(BaseModel):
+    type: Literal['CustomToolCall']
+
+
+class ToolCallUnion(RootModel[Union[FunctionToolCall, CustomToolCall]]):
+    root: Union[FunctionToolCall, CustomToolCall] = Field(..., discriminator='type')

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

@@ -0,0 +1,27 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_enum_single_value_anyof.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Literal, Union
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class ToolType(Enum):
+    function = 'function'
+
+
+class FunctionToolCall(BaseModel):
+    id: str
+    type: Literal[ToolType.function]
+
+
+class CustomToolCall(BaseModel):
+    type: Literal['CustomToolCall']
+
+
+class ToolCallUnion(RootModel[Union[FunctionToolCall, CustomToolCall]]):
+    root: Union[FunctionToolCall, CustomToolCall] = Field(..., discriminator='type')

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

@@ -0,0 +1,25 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_enum_single_value.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class ToolType(Enum):
+    function = 'function'
+
+
+class ToolBase(BaseModel):
+    type: ToolType
+
+
+class FunctionToolCall(ToolBase):
+    id: str
+
+
+class ToolCall(RootModel[FunctionToolCall]):
+    root: FunctionToolCall = Field(..., discriminator='type')

tests/data/openapi/discriminator_enum_single_value.yaml

@@ -0,0 +1,31 @@
+openapi: "3.0.0"
+components:
+  schemas:
+    ToolCall:
+      oneOf:
+        - $ref: '#/components/schemas/FunctionToolCall'
+      discriminator:
+        propertyName: type
+
+    ToolType:
+      type: string
+      enum:
+        - function
+
+    ToolBase:
+      type: object
+      properties:
+        type:
+          $ref: '#/components/schemas/ToolType'
+      required:
+        - type
+
+    FunctionToolCall:
+      allOf:
+        - $ref: '#/components/schemas/ToolBase'
+      type: object
+      properties:
+        id:
+          type: string
+      required:
+        - id

tests/data/openapi/discriminator_enum_single_value_anyof.yaml

@@ -0,0 +1,31 @@
+openapi: "3.0.0"
+components:
+  schemas:
+    ToolCallUnion:
+      anyOf:
+        - $ref: '#/components/schemas/FunctionToolCall'
+        - $ref: '#/components/schemas/CustomToolCall'
+      discriminator:
+        propertyName: type
+
+    ToolType:
+      type: string
+      enum:
+        - function
+
+    FunctionToolCall:
+      type: object
+      properties:
+        id:
+          type: string
+        type:
+          $ref: '#/components/schemas/ToolType'
+      required:
+        - id
+        - type
+
+    CustomToolCall:
+      type: object
+      properties:
+        type:
+          type: string

tests/main/openapi/test_main_openapi.py

@@ -133,6 +133,82 @@ 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_enum_single_value(output_file: Path) -> None:
+    """Single-value enum discriminator with allOf inheritance."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_enum_single_value.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "enum_single_value.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_enum_single_value_use_enum(output_file: Path) -> None:
+    """Single-value enum with allOf + --use-enum-values-in-discriminator."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_enum_single_value.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "enum_single_value_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_enum_single_value_anyof(output_file: Path) -> None:
+    """Single-value enum discriminator with anyOf - uses enum value, not model name."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_enum_single_value_anyof.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "enum_single_value_anyof.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_enum_single_value_anyof_use_enum(output_file: Path) -> None:
+    """Single-value enum with anyOf + --use-enum-values-in-discriminator."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_enum_single_value_anyof.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "enum_single_value_anyof_use_enum.py",
+        extra_args=[
+            "--target-python-version",
+            "3.10",
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+            "--use-enum-values-in-discriminator",
+        ],
+    )
+
+
 def test_main_openapi_discriminator_with_properties(output_file: Path) -> None:
     """Test OpenAPI generation with discriminator properties."""
     run_main_and_assert(