diff --git a/datamodel_code_generator/parser/base.py b/datamodel_code_generator/parser/base.py index 51953856..9a4e9a23 100644 --- a/datamodel_code_generator/parser/base.py +++ b/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, ) ) diff --git a/tests/data/expected/main/main_openapi_discriminator_enum_duplicate/output.py b/tests/data/expected/main/main_openapi_discriminator_enum_duplicate/output.py new file mode 100644 index 00000000..083fa72b --- /dev/null +++ b/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') diff --git a/tests/data/openapi/discriminator_enum_duplicate.yaml b/tests/data/openapi/discriminator_enum_duplicate.yaml new file mode 100644 index 00000000..41d0ba42 --- /dev/null +++ b/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 diff --git a/tests/test_main.py b/tests/test_main.py index b9bdfcd3..598b1d1b 100644 --- a/tests/test_main.py +++ b/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() + )