Support multiple type enum (#309)

koxudaxi · Jan 18, 2021 · aa7ca9a · aa7ca9a
1 parent dffdc4f
commit aa7ca9a
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 6 deletions.
diff --git a/datamodel_code_generator/parser/jsonschema.py b/datamodel_code_generator/parser/jsonschema.py
@@ -143,7 +143,7 @@ def validate_exclusive_maximum_and_exclusive_minimum(
     oneOf: List['JsonSchemaObject'] = []
     anyOf: List['JsonSchemaObject'] = []
     allOf: List['JsonSchemaObject'] = []
-    enum: List[str] = []
+    enum: List[Any] = []
     writeOnly: Optional[bool]
     properties: Optional[Dict[str, 'JsonSchemaObject']]
     required: List[str] = []
@@ -697,17 +697,20 @@ def parse_enum(
         enum_fields: List[DataModelFieldBase] = []
 
         for i, enum_part in enumerate(obj.enum):
-            if obj.type == 'string' or (
-                isinstance(obj.type, list) and 'string' in obj.type
-            ):
+            if obj.type == 'string' or isinstance(enum_part, str):
                 default = f"'{enum_part}'"
                 field_name = enum_part
             else:
                 default = enum_part
                 if obj.x_enum_varnames:
                     field_name = obj.x_enum_varnames[i]
                 else:
-                    field_name = f'{obj.type}_{enum_part}'
+                    prefix = (
+                        obj.type
+                        if isinstance(obj.type, str)
+                        else type(enum_part).__name__
+                    )
+                    field_name = f'{prefix}_{enum_part}'
             enum_fields.append(
                 self.data_model_field_type(
                     name=self.model_resolver.get_valid_name(field_name),

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
@@ -57,3 +57,11 @@ class AliasEnum(Enum):
     a = 1
     b = 2
     c = 3
+
+
+class MultipleTypeEnum(Enum):
+    red = 'red'
+    amber = 'amber'
+    green = 'green'
+    NoneType_None = None
+    int_42 = 42
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
@@ -59,3 +59,11 @@ class AliasEnum(Enum):
     a = 1
     b = 2
     c = 3
+
+
+class MultipleTypeEnum(Enum):
+    red = 'red'
+    amber = 'amber'
+    green = 'green'
+    NoneType_None = None
+    int_42 = 42
diff --git a/tests/data/openapi/enum_models.yaml b/tests/data/openapi/enum_models.yaml
@@ -97,4 +97,6 @@ components:
     AliasEnum:
       enum: [1,2,3]
       type: number
-      x-enum-varnames: ['a', 'b', 'c']
+      x-enum-varnames: ['a', 'b', 'c']
+    MultipleTypeEnum:
+      enum: [ "red", "amber", "green", null, 42 ]