diff --git a/json_schema_to_class.py b/json_schema_to_class.py index b8c2d4c..7c1c677 100644 --- a/json_schema_to_class.py +++ b/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: diff --git a/tests/test_json_schema_to_class.py b/tests/test_json_schema_to_class.py index 79997ec..fee6909 100644 --- a/tests/test_json_schema_to_class.py +++ b/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(): diff --git a/tests/test_schema.json b/tests/test_schema.json index 9598332..13843bf 100644 --- a/tests/test_schema.json +++ b/tests/test_schema.json @@ -24,6 +24,7 @@ "type": "object", "properties": { "lr_mode": { + "type": "string", "enum": [ "step", "cos"