Fix: Enhance handling of allOf with oneOf and anyOf references in OpenAPI generation (#2605)

Author: koxudaxi

src/datamodel_code_generator/parser/jsonschema.py

@@ -1159,9 +1159,22 @@ def _parse_all_of_item(  # noqa: PLR0913, PLR0917
     ) -> None:
         for all_of_item in obj.allOf:
             if all_of_item.ref:  # $ref
-                ref = self.model_resolver.add_ref(all_of_item.ref)
-                if ref.path not in {b.path for b in base_classes}:
-                    base_classes.append(ref)
+                ref_schema = self._load_ref_schema_object(all_of_item.ref)
+
+                if ref_schema.oneOf or ref_schema.anyOf:
+                    self.model_resolver.add(path, name, class_name=True, loaded=True)
+                    if ref_schema.anyOf:
+                        union_models.extend(
+                            d.reference for d in self.parse_any_of(name, ref_schema, path) if d.reference
+                        )
+                    if ref_schema.oneOf:
+                        union_models.extend(
+                            d.reference for d in self.parse_one_of(name, ref_schema, path) if d.reference
+                        )
+                else:
+                    ref = self.model_resolver.add_ref(all_of_item.ref)
+                    if ref.path not in {b.path for b in base_classes}:
+                        base_classes.append(ref)
             else:
                 module_name = get_module_name(name, None, treat_dot_as_module=self.treat_dot_as_module)
                 object_fields = self.parse_object_fields(

tests/data/expected/main/openapi/allof_with_anyof_ref.py

@@ -0,0 +1,48 @@
+# generated by datamodel-codegen:
+#   filename:  allof_with_anyof_ref.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Union
+
+from pydantic import BaseModel, RootModel
+
+
+class ItemType(Enum):
+    text = 'text'
+
+
+class TextItem(BaseModel):
+    itemType: ItemType
+    text: str
+
+
+class ItemType1(Enum):
+    number = 'number'
+
+
+class NumberItem(BaseModel):
+    itemType: ItemType1
+    value: int
+
+
+class Item(RootModel[Union[TextItem, NumberItem]]):
+    root: Union[TextItem, NumberItem]
+
+
+class ItemPostRequest1(BaseModel):
+    itemId: str
+
+
+class ItemPostRequest2(TextItem, ItemPostRequest1):
+    pass
+
+
+class ItemPostRequest3(NumberItem, ItemPostRequest1):
+    pass
+
+
+class ItemPostRequest(RootModel[Union[ItemPostRequest2, ItemPostRequest3]]):
+    root: Union[ItemPostRequest2, ItemPostRequest3]

tests/data/expected/main/openapi/allof_with_oneof_ref.py

@@ -0,0 +1,48 @@
+# generated by datamodel-codegen:
+#   filename:  allof_with_oneof_ref.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 UserType(Enum):
+    admin = 'admin'
+
+
+class AdminUser(BaseModel):
+    userType: Literal['admin']
+    adminLevel: int
+
+
+class UserType1(Enum):
+    regular = 'regular'
+
+
+class RegularUser(BaseModel):
+    userType: Literal['regular']
+    username: str
+
+
+class User(RootModel[Union[AdminUser, RegularUser]]):
+    root: Union[AdminUser, RegularUser] = Field(..., discriminator='userType')
+
+
+class UserPostRequest1(BaseModel):
+    userId: str
+
+
+class UserPostRequest2(AdminUser, UserPostRequest1):
+    pass
+
+
+class UserPostRequest3(RegularUser, UserPostRequest1):
+    pass
+
+
+class UserPostRequest(RootModel[Union[UserPostRequest2, UserPostRequest3]]):
+    root: Union[UserPostRequest2, UserPostRequest3]

tests/data/openapi/allof_with_anyof_ref.yaml

@@ -0,0 +1,59 @@
+# Test case for allOf referencing a schema with anyOf
+# This tests the anyOf branch in _parse_all_of_item
+
+openapi: 3.0.0
+info:
+  title: Test API
+  version: 1.0.0
+paths:
+  /items:
+    post:
+      summary: Create an item
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/ItemPostRequest'
+      responses:
+        '200':
+          description: Success
+components:
+  schemas:
+    TextItem:
+      type: object
+      required:
+        - itemType
+        - text
+      properties:
+        itemType:
+          type: string
+          enum:
+            - text
+        text:
+          type: string
+    NumberItem:
+      type: object
+      required:
+        - itemType
+        - value
+      properties:
+        itemType:
+          type: string
+          enum:
+            - number
+        value:
+          type: integer
+    Item:
+      anyOf:
+        - $ref: '#/components/schemas/TextItem'
+        - $ref: '#/components/schemas/NumberItem'
+    ItemPostRequest:
+      allOf:
+        - $ref: '#/components/schemas/Item'
+        - type: object
+          required:
+            - itemId
+          properties:
+            itemId:
+              type: string

tests/data/openapi/allof_with_oneof_ref.yaml

@@ -0,0 +1,64 @@
+# Test case for issue #1763
+# allOf referencing a schema with oneOf + discriminator
+
+openapi: 3.0.0
+info:
+  title: Test API
+  version: 1.0.0
+paths:
+  /users:
+    post:
+      summary: Create a user
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/UserPostRequest'
+      responses:
+        '200':
+          description: Success
+components:
+  schemas:
+    AdminUser:
+      type: object
+      required:
+        - userType
+        - adminLevel
+      properties:
+        userType:
+          type: string
+          enum:
+            - admin
+        adminLevel:
+          type: integer
+    RegularUser:
+      type: object
+      required:
+        - userType
+        - username
+      properties:
+        userType:
+          type: string
+          enum:
+            - regular
+        username:
+          type: string
+    User:
+      oneOf:
+        - $ref: '#/components/schemas/AdminUser'
+        - $ref: '#/components/schemas/RegularUser'
+      discriminator:
+        propertyName: userType
+        mapping:
+          admin: '#/components/schemas/AdminUser'
+          regular: '#/components/schemas/RegularUser'
+    UserPostRequest:
+      allOf:
+        - $ref: '#/components/schemas/User'
+        - type: object
+          required:
+            - userId
+          properties:
+            userId:
+              type: string

tests/main/openapi/test_main_openapi.py

@@ -135,6 +135,44 @@ def test_main_openapi_discriminator_allof_no_subtypes(output_file: Path) -> None
     )
 
 
+def test_main_openapi_allof_with_oneof_ref(output_file: Path) -> None:
+    """Test OpenAPI generation with allOf referencing a oneOf schema.
+
+    This tests the case where allOf combines a $ref to a schema with oneOf/discriminator
+    and additional properties. Regression test for issue #1763.
+    """
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "allof_with_oneof_ref.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "allof_with_oneof_ref.py",
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+        ],
+    )
+
+
+def test_main_openapi_allof_with_anyof_ref(output_file: Path) -> None:
+    """Test OpenAPI generation with allOf referencing an anyOf schema.
+
+    This tests the case where allOf combines a $ref to a schema with anyOf
+    and additional properties.
+    """
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "allof_with_anyof_ref.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "allof_with_anyof_ref.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(