From af9f495334e3c9de0f333555f827c3bac7b657d3 Mon Sep 17 00:00:00 2001 From: Philip Bjorge Date: Wed, 5 May 2021 10:20:37 -0700 Subject: [PATCH] Handle nested nullable enums (#425) --- datamodel_code_generator/parser/jsonschema.py | 5 ++++- .../main_openapi_enum_models_all/output.py | 6 ++++++ .../output.py | 6 ++++++ .../main_openapi_enum_models_one/output.py | 21 +++++++++++++++++++ .../output_py36.py | 21 +++++++++++++++++++ .../output_py37.py | 21 +++++++++++++++++++ tests/data/openapi/enum_models.yaml | 17 +++++++++++++++ 7 files changed, 96 insertions(+), 1 deletion(-) diff --git a/datamodel_code_generator/parser/jsonschema.py b/datamodel_code_generator/parser/jsonschema.py index ffc27af6f..739874907 100644 --- a/datamodel_code_generator/parser/jsonschema.py +++ b/datamodel_code_generator/parser/jsonschema.py @@ -630,7 +630,10 @@ def parse_item( return self.data_type_manager.get_data_type(Types.object) elif item.enum: if self.should_parse_enum_as_literal(item): - return self.data_type(literals=item.enum) + enum_literals = item.enum + if item.nullable: + enum_literals = [i for i in item.enum if i is not None] + return self.data_type(literals=enum_literals) return self.parse_enum( name, item, get_special_path('enum', path), singular_name=singular_name ) diff --git a/tests/data/expected/main/main_openapi_enum_models_all/output.py b/tests/data/expected/main/main_openapi_enum_models_all/output.py index 2bcd255fe..9aa24b9bc 100644 --- a/tests/data/expected/main/main_openapi_enum_models_all/output.py +++ b/tests/data/expected/main/main_openapi_enum_models_all/output.py @@ -67,6 +67,12 @@ class ArrayEnum(BaseModel): __root__: List[Union[Literal['cat'], Literal['dog']]] +class NestedNullableEnum(BaseModel): + nested_version: Optional[ + Literal['RC1', 'RC1N', 'RC2', 'RC2N', 'RC3', 'RC4'] + ] = Field('RC1', description='nullable enum', example='RC2') + + class VersionEnum(Enum): RC1 = 'RC1' RC1N = 'RC1N' diff --git a/tests/data/expected/main/main_openapi_enum_models_as_literal_py37/output.py b/tests/data/expected/main/main_openapi_enum_models_as_literal_py37/output.py index a8108f042..11c797c3f 100644 --- a/tests/data/expected/main/main_openapi_enum_models_as_literal_py37/output.py +++ b/tests/data/expected/main/main_openapi_enum_models_as_literal_py37/output.py @@ -68,6 +68,12 @@ class ArrayEnum(BaseModel): __root__: List[Union[Literal['cat'], Literal['dog']]] +class NestedNullableEnum(BaseModel): + nested_version: Optional[ + Literal['RC1', 'RC1N', 'RC2', 'RC2N', 'RC3', 'RC4'] + ] = Field('RC1', description='nullable enum', example='RC2') + + class VersionEnum(Enum): RC1 = 'RC1' RC1N = 'RC1N' diff --git a/tests/data/expected/main/main_openapi_enum_models_one/output.py b/tests/data/expected/main/main_openapi_enum_models_one/output.py index f6459d721..baf435faa 100644 --- a/tests/data/expected/main/main_openapi_enum_models_one/output.py +++ b/tests/data/expected/main/main_openapi_enum_models_one/output.py @@ -82,6 +82,27 @@ class ArrayEnum(BaseModel): __root__: List[Union[Literal['cat'], Literal['dog']]] +class NestedVersionEnum(Enum): + RC1 = 'RC1' + RC1N = 'RC1N' + RC2 = 'RC2' + RC2N = 'RC2N' + RC3 = 'RC3' + RC4 = 'RC4' + + +class NestedVersion(BaseModel): + __root__: Optional[NestedVersionEnum] = Field( + 'RC1', description='nullable enum', example='RC2' + ) + + +class NestedNullableEnum(BaseModel): + nested_version: Optional[NestedVersion] = Field( + 'RC1', description='nullable enum', example='RC2' + ) + + class VersionEnum(Enum): RC1 = 'RC1' RC1N = 'RC1N' diff --git a/tests/data/expected/parser/openapi/openapi_parser_parse_enum_models/output_py36.py b/tests/data/expected/parser/openapi/openapi_parser_parse_enum_models/output_py36.py index 2f1482600..b1c9e9e09 100644 --- a/tests/data/expected/parser/openapi/openapi_parser_parse_enum_models/output_py36.py +++ b/tests/data/expected/parser/openapi/openapi_parser_parse_enum_models/output_py36.py @@ -88,6 +88,27 @@ class ArrayEnum(BaseModel): __root__: List[Union['ArrayEnumEnum', 'ArrayEnumEnum1']] +class NestedVersionEnum(Enum): + RC1 = 'RC1' + RC1N = 'RC1N' + RC2 = 'RC2' + RC2N = 'RC2N' + RC3 = 'RC3' + RC4 = 'RC4' + + +class NestedVersion(BaseModel): + __root__: Optional['NestedVersionEnum'] = Field( + 'RC1', description='nullable enum', example='RC2' + ) + + +class NestedNullableEnum(BaseModel): + nested_version: Optional['NestedVersion'] = Field( + 'RC1', description='nullable enum', example='RC2' + ) + + class VersionEnum(Enum): RC1 = 'RC1' RC1N = 'RC1N' diff --git a/tests/data/expected/parser/openapi/openapi_parser_parse_enum_models/output_py37.py b/tests/data/expected/parser/openapi/openapi_parser_parse_enum_models/output_py37.py index ad53957ef..a7c7db10f 100644 --- a/tests/data/expected/parser/openapi/openapi_parser_parse_enum_models/output_py37.py +++ b/tests/data/expected/parser/openapi/openapi_parser_parse_enum_models/output_py37.py @@ -90,6 +90,27 @@ class ArrayEnum(BaseModel): __root__: List[Union[ArrayEnumEnum, ArrayEnumEnum1]] +class NestedVersionEnum(Enum): + RC1 = 'RC1' + RC1N = 'RC1N' + RC2 = 'RC2' + RC2N = 'RC2N' + RC3 = 'RC3' + RC4 = 'RC4' + + +class NestedVersion(BaseModel): + __root__: Optional[NestedVersionEnum] = Field( + 'RC1', description='nullable enum', example='RC2' + ) + + +class NestedNullableEnum(BaseModel): + nested_version: Optional[NestedVersion] = Field( + 'RC1', description='nullable enum', example='RC2' + ) + + class VersionEnum(Enum): RC1 = 'RC1' RC1N = 'RC1N' diff --git a/tests/data/openapi/enum_models.yaml b/tests/data/openapi/enum_models.yaml index dc7ced2f3..2892912fb 100644 --- a/tests/data/openapi/enum_models.yaml +++ b/tests/data/openapi/enum_models.yaml @@ -112,6 +112,23 @@ components: { enum: [ "cat" ] }, { enum: [ "dog"]} ] + nestedNullableEnum: + type: object + properties: + nested_version: + type: string + nullable: true + default: RC1 + description: nullable enum + example: RC2 + enum: + - RC1 + - RC1N + - RC2 + - RC2N + - RC3 + - RC4 + - null version: type: string nullable: true