Support allOf polymorphism with discriminator in OpenAPI (#2530)

* Support allOf polymorphism in openapi

* move logic to openapi.py

* Apply suggestions from code review

Co-authored-by: Antonio Spadaro <ilovelinux@users.noreply.github.com>

* refactor parse_object_fields for code reuse

* remove target python version from test

* Refactor discriminator handling to improve polymorphism support

* Enhance discriminator handling for allOf polymorphism support

* Add test for OpenAPI generation with discriminator and no allOf subtypes

* Update src/datamodel_code_generator/parser/openapi.py

---------

Co-authored-by: Antonio Spadaro <ilovelinux@users.noreply.github.com>
Co-authored-by: Koudai Aono <koxudaxi@gmail.com>

Author: 3 people

src/datamodel_code_generator/parser/openapi.py

@@ -344,6 +344,8 @@ def __init__(  # noqa: PLR0913
         )
         self.open_api_scopes: list[OpenAPIScope] = openapi_scopes or [OpenAPIScope.Schemas]
         self.include_path_parameters: bool = include_path_parameters
+        self._discriminator_schemas: dict[str, dict[str, Any]] = {}
+        self._discriminator_subtypes: dict[str, list[str]] = defaultdict(list)
 
     def get_ref_model(self, ref: str) -> dict[str, Any]:
         """Resolve a reference to its model definition."""
@@ -362,6 +364,49 @@ def get_data_type(self, obj: JsonSchemaObject) -> DataType:
 
         return super().get_data_type(obj)
 
+    def _get_discriminator_union_type(self, ref: str) -> DataType | None:
+        """Create a union type for discriminator subtypes if available."""
+        subtypes = self._discriminator_subtypes.get(ref, [])
+        if not subtypes:
+            return None
+        refs = map(self.model_resolver.add_ref, subtypes)
+        return self.data_type(data_types=[self.data_type(reference=r) for r in refs])
+
+    def get_ref_data_type(self, ref: str) -> DataType:
+        """Get data type for a reference, handling discriminator polymorphism."""
+        if ref in self._discriminator_schemas and (union_type := self._get_discriminator_union_type(ref)):
+            return union_type
+        return super().get_ref_data_type(ref)
+
+    def parse_object_fields(
+        self,
+        obj: JsonSchemaObject,
+        path: list[str],
+        module_name: Optional[str] = None,  # noqa: UP045
+    ) -> list[DataModelFieldBase]:
+        """Parse object fields, adding discriminator info for allOf polymorphism."""
+        fields = super().parse_object_fields(obj, path, module_name)
+        properties = obj.properties or {}
+
+        result_fields: list[DataModelFieldBase] = []
+        for field_obj in fields:
+            field = properties.get(field_obj.original_name)
+
+            if (
+                isinstance(field, JsonSchemaObject)
+                and field.ref
+                and (discriminator := self._discriminator_schemas.get(field.ref))
+            ):
+                new_field_type = self._get_discriminator_union_type(field.ref) or field_obj.data_type
+                field_obj = self.data_model_field_type(**{  # noqa: PLW2901
+                    **field_obj.__dict__,
+                    "data_type": new_field_type,
+                    "extras": {**field_obj.extras, "discriminator": discriminator},
+                })
+            result_fields.append(field_obj)
+
+        return result_fields
+
     def resolve_object(self, obj: ReferenceObject | BaseModelT, object_type: type[BaseModelT]) -> BaseModelT:
         """Resolve a reference object to its actual type or return the object as-is."""
         if isinstance(obj, ReferenceObject):
@@ -651,6 +696,7 @@ def parse_raw(self) -> None:  # noqa: PLR0912, PLR0915
 
             specification: dict[str, Any] = load_yaml_dict(source.text)
             self.raw_obj = specification
+            self._collect_discriminator_schemas()
             schemas: dict[str, Any] = specification.get("components", {}).get("schemas", {})
             security: list[dict[str, list[str]]] | None = specification.get("security")
             if OpenAPIScope.Schemas in self.open_api_scopes:
@@ -727,3 +773,25 @@ def parse_raw(self) -> None:  # noqa: PLR0912, PLR0915
                         )
 
         self._resolve_unparsed_json_pointer()
+
+    def _collect_discriminator_schemas(self) -> None:
+        """Collect schemas with discriminators but no oneOf/anyOf, and find their subtypes."""
+        schemas: dict[str, Any] = self.raw_obj.get("components", {}).get("schemas", {})
+
+        for schema_name, schema in schemas.items():
+            discriminator = schema.get("discriminator")
+            if not discriminator:
+                continue
+
+            if schema.get("oneOf") or schema.get("anyOf"):
+                continue
+
+            ref = f"#/components/schemas/{schema_name}"
+            self._discriminator_schemas[ref] = discriminator
+
+        for schema_name, schema in schemas.items():
+            for all_of_item in schema.get("allOf", []):
+                ref_in_allof = all_of_item.get("$ref")
+                if ref_in_allof and ref_in_allof in self._discriminator_schemas:
+                    subtype_ref = f"#/components/schemas/{schema_name}"
+                    self._discriminator_subtypes[ref_in_allof].append(subtype_ref)

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

@@ -0,0 +1,32 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_allof.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+
+class Pet(BaseModel):
+    pet_type: Annotated[str, Field(alias='petType')]
+
+
+class Cat(Pet):
+    name: str | None = None
+    pet_type: Literal['cat'] = Field(..., alias='petType')
+
+
+class Dog(Pet):
+    bark: str | None = None
+    pet_type: Literal['dog'] = Field(..., alias='petType')
+
+
+class Lizard(Pet):
+    loves_rocks: Annotated[bool | None, Field(alias='lovesRocks')] = None
+    pet_type: Literal['lizard'] = Field(..., alias='petType')
+
+
+class PetContainer(BaseModel):
+    pet: Annotated[Cat | Dog | Lizard, Field(discriminator='pet_type')]

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

@@ -0,0 +1,25 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_allof_no_subtypes.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class BaseItem(BaseModel):
+    itemType: str
+
+
+class FooItem(BaseModel):
+    fooValue: Optional[str] = None
+
+
+class BarItem(BaseModel):
+    barValue: Optional[int] = None
+
+
+class ItemContainer(BaseModel):
+    item: BaseItem = Field(..., discriminator='itemType')

tests/data/openapi/discriminator_allof.yaml

@@ -0,0 +1,62 @@
+# Example from https://spec.openapis.org/oas/v3.1.1.html#examples-1
+# This tests discriminator without oneOf/anyOf, where subtypes use allOf
+
+openapi: 3.1.0
+info:
+  title: Test
+  description: "Test API"
+  version: 0.0.0
+paths:
+  /pet:
+    get:
+      responses:
+        '200':
+          description: "Pet"
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/PetContainer"
+components:
+  schemas:
+    PetContainer:
+      type: object
+      required:
+        - pet
+      properties:
+        pet:
+          $ref: "#/components/schemas/Pet"
+    Pet:
+      type: object
+      required:
+        - petType
+      properties:
+        petType:
+          type: string
+      discriminator:
+        propertyName: petType
+        mapping:
+          cat: "#/components/schemas/Cat"
+          dog: "#/components/schemas/Dog"
+          lizard: "#/components/schemas/Lizard"
+    Cat:
+      allOf:
+        - $ref: "#/components/schemas/Pet"
+        - type: object
+          properties:
+            name:
+              type: string
+    Dog:
+      allOf:
+        - $ref: "#/components/schemas/Pet"
+        - type: object
+          properties:
+            bark:
+              type: string
+    Lizard:
+      allOf:
+        - $ref: "#/components/schemas/Pet"
+        - type: object
+          properties:
+            lovesRocks:
+              type: boolean
+

tests/data/openapi/discriminator_allof_no_subtypes.yaml

@@ -0,0 +1,50 @@
+# Tests discriminator without any allOf subtypes
+# This tests the edge case where a schema has a discriminator but nothing inherits from it
+
+openapi: 3.1.0
+info:
+  title: Test
+  description: "Test API"
+  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"
+          bar: "#/components/schemas/BarItem"
+    # These schemas exist but don't use allOf to inherit from BaseItem
+    FooItem:
+      type: object
+      properties:
+        fooValue:
+          type: string
+    BarItem:
+      type: object
+      properties:
+        barValue:
+          type: integer

tests/main/openapi/test_main_openapi.py

@@ -97,6 +97,44 @@ def test_main_openapi_discriminator_with_properties(output_file: Path) -> None:
     )
 
 
+def test_main_openapi_discriminator_allof(output_file: Path) -> None:
+    """Test OpenAPI generation with allOf discriminator polymorphism."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_allof.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "allof.py",
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+            "--snake-case-field",
+            "--use-annotated",
+            "--use-union-operator",
+            "--collapse-root-models",
+        ],
+    )
+
+
+def test_main_openapi_discriminator_allof_no_subtypes(output_file: Path) -> None:
+    """Test OpenAPI generation with discriminator but no allOf subtypes.
+
+    This tests the edge case where a schema has a discriminator but nothing
+    inherits from it using allOf.
+    """
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_allof_no_subtypes.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "allof_no_subtypes.py",
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+        ],
+    )
+
+
 def test_main_pydantic_basemodel(output_file: Path) -> None:
     """Test OpenAPI generation with Pydantic BaseModel output."""
     run_main_and_assert(