Fix discriminator with allOf without Literal type for Pydantic v2 (#2722)

Author: koxudaxi

src/datamodel_code_generator/parser/openapi.py

@@ -395,9 +395,53 @@ def get_data_type(self, obj: JsonSchemaObject) -> DataType:
 
         return super().get_data_type(obj)
 
+    def _normalize_discriminator_mapping_ref(self, mapping_value: str) -> str:  # noqa: PLR6301
+        """Normalize a discriminator mapping value to a full $ref path.
+
+        Per OpenAPI spec, mapping values can be either:
+        - Full refs: "#/components/schemas/Pet" or "./other.yaml#/components/schemas/Pet"
+        - Short names: "Pet" or "Pet.V1" (relative to #/components/schemas/)
+        - Relative paths: "schemas/Pet" or "./other.yaml"
+
+        Values containing "/" or "#" are treated as paths/refs and passed through.
+        All other values (including those with dots like "Pet.V1") are treated as
+        short schema names and normalized to full refs.
+
+        Note: Bare file references without path separators (e.g., "other.yaml") will be
+        treated as schema names. Use "./other.yaml" format for file references.
+
+        Note: This could be a staticmethod, but @snooper_to_methods() decorator
+        converts staticmethods to regular functions when pysnooper is installed.
+        """
+        if "/" in mapping_value or "#" in mapping_value:
+            return mapping_value
+        return f"#/components/schemas/{mapping_value}"
+
+    def _normalize_discriminator(self, discriminator: dict[str, Any]) -> dict[str, Any]:
+        """Return a copy of the discriminator dict with normalized mapping refs."""
+        result = discriminator.copy()
+        mapping = discriminator.get("mapping")
+        if mapping:
+            result["mapping"] = {
+                k: self._normalize_discriminator_mapping_ref(v) for k, v in mapping.items() if isinstance(v, str)
+            }
+        return result
+
     def _get_discriminator_union_type(self, ref: str) -> DataType | None:
-        """Create a union type for discriminator subtypes if available."""
+        """Create a union type for discriminator subtypes if available.
+
+        First tries to use allOf subtypes. If none found, falls back to using
+        the discriminator mapping to create the union type. This handles cases
+        where schemas don't use allOf inheritance but have explicit discriminator mappings.
+        """
         subtypes = self._discriminator_subtypes.get(ref, [])
+        if not subtypes:
+            discriminator = self._discriminator_schemas[ref]
+            mapping = discriminator.get("mapping", {})
+            if mapping:
+                subtypes = [
+                    self._normalize_discriminator_mapping_ref(v) for v in mapping.values() if isinstance(v, str)
+                ]
         if not subtypes:
             return None
         refs = map(self.model_resolver.add_ref, subtypes)
@@ -430,10 +474,11 @@ def parse_object_fields(
                 and (discriminator := self._discriminator_schemas.get(field.ref))
             ):
                 new_field_type = self._get_discriminator_union_type(field.ref) or field_obj.data_type
+                normalized_discriminator = self._normalize_discriminator(discriminator)
                 field_obj = self.data_model_field_type(**{  # noqa: PLW2901
                     **field_obj.__dict__,
                     "data_type": new_field_type,
-                    "extras": {**field_obj.extras, "discriminator": discriminator},
+                    "extras": {**field_obj.extras, "discriminator": normalized_discriminator},
                 })
             result_fields.append(field_obj)
 

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

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+from typing import Literal
+
 from pydantic import BaseModel, Field
 
 
@@ -13,11 +15,13 @@ class BaseItem(BaseModel):
 
 class FooItem(BaseModel):
     fooValue: str | None = None
+    itemType: Literal['foo']
 
 
 class BarItem(BaseModel):
     barValue: int | None = None
+    itemType: Literal['bar']
 
 
 class ItemContainer(BaseModel):
-    item: BaseItem = Field(..., discriminator='itemType')
+    item: FooItem | BarItem = Field(..., discriminator='itemType')

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

@@ -0,0 +1,27 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_no_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['FooItem']
+
+
+class BarItem(BaseItem):
+    barValue: int | None = None
+    itemType: Literal['BarItem']
+
+
+class ItemContainer(BaseModel):
+    item: FooItem | BarItem = Field(..., discriminator='itemType')

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

@@ -0,0 +1,19 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_no_mapping_no_subtypes.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class BaseItem(BaseModel):
+    itemType: str
+
+
+class FooItem(BaseModel):
+    fooValue: str | None = None
+
+
+class ItemContainer(BaseModel):
+    item: BaseItem = Field(..., discriminator='itemType')

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

@@ -0,0 +1,27 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_short_mapping_names.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(BaseModel):
+    fooValue: str | None = None
+    itemType: Literal['foo']
+
+
+class BarItem(BaseModel):
+    barValue: int | None = None
+    itemType: Literal['bar']
+
+
+class ItemContainer(BaseModel):
+    item: FooItem | BarItem = Field(..., discriminator='itemType')

tests/data/openapi/discriminator_no_mapping.yaml

@@ -0,0 +1,47 @@
+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
+    FooItem:
+      allOf:
+        - $ref: "#/components/schemas/BaseItem"
+        - type: object
+          properties:
+            fooValue:
+              type: string
+    BarItem:
+      allOf:
+        - $ref: "#/components/schemas/BaseItem"
+        - type: object
+          properties:
+            barValue:
+              type: integer

tests/data/openapi/discriminator_no_mapping_no_subtypes.yaml

@@ -0,0 +1,38 @@
+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
+    FooItem:
+      type: object
+      properties:
+        fooValue:
+          type: string

tests/data/openapi/discriminator_short_mapping_names.yaml

@@ -0,0 +1,50 @@
+# Tests discriminator with short mapping names (not full $ref paths)
+# Per OpenAPI spec, mapping values can be short names like "FooItem" instead of "#/components/schemas/FooItem"
+
+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:
+          # Short names without full $ref paths
+          foo: FooItem
+          bar: BarItem
+    FooItem:
+      type: object
+      properties:
+        fooValue:
+          type: string
+    BarItem:
+      type: object
+      properties:
+        barValue:
+          type: integer

tests/main/openapi/test_main_openapi.py

@@ -271,6 +271,64 @@ def test_main_openapi_discriminator_allof_no_subtypes(output_file: Path) -> None
     )
 
 
+def test_main_openapi_discriminator_short_mapping_names(output_file: Path) -> None:
+    """Test OpenAPI generation with discriminator using short mapping names.
+
+    Per OpenAPI spec, mapping values can be short names like "FooItem" instead
+    of full refs like "#/components/schemas/FooItem". This tests that short
+    names are normalized correctly.
+    """
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_short_mapping_names.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "short_mapping_names.py",
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+        ],
+    )
+
+
+def test_main_openapi_discriminator_no_mapping(output_file: Path) -> None:
+    """Test OpenAPI generation with discriminator without mapping.
+
+    This tests the case where a discriminator has only propertyName but no mapping.
+    The subtypes are discovered via allOf inheritance.
+    """
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_no_mapping.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "no_mapping.py",
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+        ],
+    )
+
+
+def test_main_openapi_discriminator_no_mapping_no_subtypes(output_file: Path) -> None:
+    """Test OpenAPI generation with discriminator without mapping and no allOf subtypes.
+
+    This tests the edge case where a discriminator has no mapping and no schemas
+    inherit from it using allOf.
+    """
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "discriminator_no_mapping_no_subtypes.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file=EXPECTED_OPENAPI_PATH / "discriminator" / "no_mapping_no_subtypes.py",
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+        ],
+    )
+
+
 def test_main_openapi_allof_with_oneof_ref(output_file: Path) -> None:
     """Test OpenAPI generation with allOf referencing a oneOf schema.