Handle nested nullable enums (#425)

koxudaxi · May 5, 2021 · af9f495 · af9f495
1 parent fe505d9
commit af9f495
Show file tree

Hide file tree

Showing 7 changed files with 96 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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