Make discriminators work with multiple keys pointing to the same schema (#1885)

* Make discriminators work with multiple keys pointing to the same schema

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Koudai Aono <koxudaxi@gmail.com>

datamodel_code_generator/parser/base.py

@@ -756,7 +756,7 @@ class Parser(ABC):
                         (pydantic_model.BaseModel, pydantic_model_v2.BaseModel),
                     ):
                         continue  # pragma: no cover
-                    type_name = None
+                    type_names = []
                     if mapping:
                         for name, path in mapping.items():
                             if (
@@ -765,10 +765,10 @@ class Parser(ABC):
                             ):
                                 # TODO: support external reference
                                 continue
-                            type_name = name
+                            type_names.append(name)
                     else:
-                        type_name = discriminator_model.path.split('/')[-1]
-                    if not type_name:  # pragma: no cover
+                        type_names = [discriminator_model.path.split('/')[-1]]
+                    if not type_names:  # pragma: no cover
                         raise RuntimeError(
                             f'Discriminator type is not found. {data_type.reference.path}'
                         )
@@ -780,7 +780,11 @@ class Parser(ABC):
                         ) != property_name:
                             continue
                         literals = discriminator_field.data_type.literals
-                        if len(literals) == 1 and literals[0] == type_name:
+                        if (
+                            len(literals) == 1 and literals[0] == type_names[0]
+                            if type_names
+                            else None
+                        ):
                             has_one_literal = True
                             continue
                         for (
@@ -789,7 +793,7 @@ class Parser(ABC):
                             if field_data_type.reference:  # pragma: no cover
                                 field_data_type.remove_reference()
                         discriminator_field.data_type = self.data_type(
-                            literals=[type_name]
+                            literals=type_names
                         )
                         discriminator_field.data_type.parent = discriminator_field
                         discriminator_field.required = True
@@ -799,7 +803,7 @@ class Parser(ABC):
                         discriminator_model.fields.append(
                             self.data_model_field_type(
                                 name=property_name,
-                                data_type=self.data_type(literals=[type_name]),
+                                data_type=self.data_type(literals=type_names),
                                 required=True,
                             )
                         )

tests/data/expected/main/main_openapi_discriminator_enum_duplicate/output.py

@@ -0,0 +1,37 @@
+# generated by datamodel-codegen:
+#   filename:  discriminator_enum_duplicate.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Literal, Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class Cat(BaseModel):
+    pet_type: Literal['cat'] = Field(..., title='Pet Type')
+    meows: int = Field(..., title='Meows')
+
+
+class Dog(BaseModel):
+    pet_type: Literal['dog'] = Field(..., title='Pet Type')
+    barks: float = Field(..., title='Barks')
+
+
+class PetType(Enum):
+    reptile = 'reptile'
+    lizard = 'lizard'
+
+
+class Lizard(BaseModel):
+    pet_type: Literal['lizard', 'reptile'] = Field(..., title='Pet Type')
+    scales: bool = Field(..., title='Scales')
+
+
+class Animal(BaseModel):
+    pet: Optional[Union[Cat, Dog, Lizard]] = Field(
+        None, discriminator='pet_type', title='Pet'
+    )
+    n: Optional[int] = Field(None, title='N')

tests/data/openapi/discriminator_enum_duplicate.yaml

@@ -0,0 +1,64 @@
+# Example from https://docs.pydantic.dev/latest/concepts/unions/#discriminated-unions
+openapi: 3.1.0
+components:
+  schemas:
+    Cat:
+      properties:
+        pet_type:
+          const: "cat"
+          title: "Pet Type"
+        meows:
+          title: Meows
+          type: integer
+      required:
+        - pet_type
+        - meows
+      title: Cat
+      type: object
+    Dog:
+      properties:
+        pet_type:
+          const: "dog"
+          title: "Pet Type"
+        barks:
+          title: Barks
+          type: number
+      required:
+        - pet_type
+        - barks
+      title: Dog
+      type: object
+    Lizard:
+      properties:
+        pet_type:
+          enum:
+            - reptile
+            - lizard
+          title: Pet Type
+          type: string
+        scales:
+          title: Scales
+          type: boolean
+      required:
+        - pet_type
+        - scales
+      title: Lizard
+      type: object
+    Animal:
+      properties:
+        pet:
+          discriminator:
+            mapping:
+              cat: '#/components/schemas/Cat'
+              dog: '#/components/schemas/Dog'
+              lizard: '#/components/schemas/Lizard'
+              reptile: '#/components/schemas/Lizard'
+            propertyName: pet_type
+          oneOf:
+            - $ref: '#/components/schemas/Cat'
+            - $ref: '#/components/schemas/Dog'
+            - $ref: '#/components/schemas/Lizard'
+          title: Pet
+        'n':
+          title: 'N'
+          type: integer

tests/test_main.py

@@ -6638,3 +6638,36 @@ def test_main_openapi_discriminator_enum():
                 EXPECTED_MAIN_PATH / 'main_openapi_discriminator_enum' / 'output.py'
             ).read_text()
         )
+
+
+@freeze_time('2019-07-26')
+@pytest.mark.skipif(
+    black.__version__.split('.')[0] == '19',
+    reason="Installed black doesn't support the old style",
+)
+def test_main_openapi_discriminator_enum_duplicate():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(OPEN_API_DATA_PATH / 'discriminator_enum_duplicate.yaml'),
+                '--output',
+                str(output_file),
+                '--target-python-version',
+                '3.10',
+                '--output-model-type',
+                'pydantic_v2.BaseModel',
+                '--input-file-type',
+                'openapi',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (
+                EXPECTED_MAIN_PATH
+                / 'main_openapi_discriminator_enum_duplicate'
+                / 'output.py'
+            ).read_text()
+        )