Add Enum to Schema generation

json_schema_to_class.py

@@ -190,6 +190,43 @@ def to_class_code(self, level: int = 0, schema: dict = None) -> str:
         return indent_class(code=code, level=level)
 
 
+class Enum(Array):
+    def __init__(self, name: str, items: type, default: Any = None):
+        super().__init__(name=name, items=items, default=default)
+        self.enumlist = items
+
+    def to_init_code(self) -> str:
+        Array.use_enum = True
+
+        if self.default is None:
+            return '{if_spaces}self.{name} = None\n{if_spaces}if "{name}" in values:\n{statement_spaces}self.{name} = self.{class_name}(values.get("{name}"))'.format(
+                if_spaces=spaces(2),
+                statement_spaces=spaces(3),
+                name=self.name,
+                class_name=self.class_name(),
+                default=self.default
+            )
+        return '{spaces}self.{name} = self.{class_name}(values.get("{name}", "{default}"))'.format(
+            spaces=spaces(2),
+            name=self.name,
+            class_name=self.class_name(),
+            default=self.default
+        )
+
+    def to_class_code(self, level: int = 0, schema: dict = None) -> str:
+        result = [f'class {self.class_name()}(Enum):']
+        self.generate_schema_code(result, schema)
+        for value in self.enumlist:
+            if value == "NONE":
+                result.append(f'{spaces(1)}{value.lower()} = "{value}"')
+            else:
+                result.append(f'{spaces(1)}{value.capitalize()} = "{value}"')
+        result.append(f'\n{spaces(1)}def __str__(self):\n{spaces(2)}return str(self.value)')
+        self.generate_validate_code(result, schema)
+        code = Config.line_break.join(result)
+        return indent_class(code=code, level=level)
+
+
 class Parser:
     def __init__(self):
         self.definitions: Dict[str, Item] = {}
@@ -211,19 +248,20 @@ def parse_definition(self, name: str, schema: dict) -> Item:
 
         if 'type' in schema:
             item_type = schema['type']
-            if item_type == 'object' and 'properties' in schema:
+            if 'enum' in schema:
+                Array.use_enum = True
+                enum_list = schema['enum']
+                assert len(enum_list) > 0, "Enum List is Empty"
+                first = enum_list[0]
+                assert all(type(first) == type(item) for item in enum_list), "Items in Enum List with Different Types"
+                assert type(first) in {int, float, str}, "Enum Type is not int, float or string"
+                return Enum(name=name, items=enum_list, default=default)
+            elif item_type == 'object' and 'properties' in schema:
                 return self.parse_object(name=name, schema=schema)
             elif item_type == 'array':
                 return self.parse_array(name=name, schema=schema)
             else:
                 return Basic(name=name, typename=Basic.TYPE_MAP[item_type], default=default)
-        elif 'enum' in schema:
-            enum_list = schema['enum']
-            assert len(enum_list) > 0, "Enum List is Empty"
-            first = enum_list[0]
-            assert all(type(first) == type(item) for item in enum_list), "Items in Enum List with Different Types"
-            assert type(first) in {int, float, str}, "Enum Type is not int, float or string"
-            return Basic(name=name, typename=type(first), default=default)
         elif '$ref' in schema:
             path: str = schema['$ref']
             class_type = path.split('/')[-1]
@@ -241,6 +279,7 @@ def parse(self, schema: dict):
 
     def generate(self, schema: dict) -> str:
         Array.use_list = False
+        Array.use_enum = False
         result = []
         for _, definition in self.definitions.items():
             result.append(definition.to_class_code(level=0))
@@ -251,6 +290,8 @@ def generate(self, schema: dict) -> str:
         headers = []
         if Array.use_list:
             headers += ['from typing import List']
+        if Array.use_enum:
+            headers += ['from enum import Enum']
         if Config.generate_validate_code:
             headers = ['import json'] + headers + ['', 'import jsonschema']
         if Config.generated_warning:

tests/test_json_schema_to_class.py

@@ -2,6 +2,7 @@
 import json
 import unittest
 from pathlib import Path
+from enum import Enum
 
 import json_schema_to_class
 
@@ -70,19 +71,21 @@ def test_array(self):
         self.assertTrue(True)
 
     def test_enum(self):
-        def parse_empty_definition():
-            parser = json_schema_to_class.Parser()
-            parser.parse(schema={
-                'title': 'days',
-                'type': 'object',
-                'properties': {
-                    'mode': {
-                        'enum': ['cosine', 'linear', 0]
-                    }
+        parser = json_schema_to_class.Parser()
+        parser.parse(schema={
+            'title': 'days',
+            'type': 'object',
+            'properties': {
+                'mode': {
+                    'type': "string",
+                    'enum': ['cosine', 'linear']
                 }
-            })
+            }
+        })
 
-        self.assertRaises(AssertionError, parse_empty_definition)
+        code = parser.root.to_class_code()
+        exec(code)
+        self.assertTrue(True)
 
     def test_parse(self):
         def parse_empty_definition():

tests/test_schema.json

@@ -24,6 +24,7 @@
       "type": "object",
       "properties": {
         "lr_mode": {
+          "type": "string",
           "enum": [
             "step",
             "cos"