Fix non-string OpenAPI discriminator literals (#3070)

* Fix non-string OpenAPI discriminator literals

* Tighten discriminator value types

Author: koxudaxi

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,25 +1515,25 @@ 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:
         """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 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,15 +1701,15 @@ 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
                         imports.append(discriminator_field.imports)
                         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

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')

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')

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')

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')

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

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

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