From 80bd7abaaafd6cb07cbfa41aa3a5a54d6d1c36a0 Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Sat, 29 May 2021 00:42:05 +0900 Subject: [PATCH] Support customTypePath (#436) * Support customTypePath * Fix unittest --- .../model/pydantic/__init__.py | 1 + .../model/pydantic/base_model.py | 4 ++ datamodel_code_generator/parser/jsonschema.py | 9 +++++ datamodel_code_generator/types.py | 10 +++++ .../output.py | 36 ++++++++++++++++++ tests/data/jsonschema/custom_type_path.json | 37 +++++++++++++++++++ tests/test_main.py | 30 +++++++++++++++ 7 files changed, 127 insertions(+) create mode 100644 tests/data/expected/main/main_jsonschema_custom_type_path/output.py create mode 100644 tests/data/jsonschema/custom_type_path.json diff --git a/datamodel_code_generator/model/pydantic/__init__.py b/datamodel_code_generator/model/pydantic/__init__.py index a498ed822..e9cbe2dd1 100644 --- a/datamodel_code_generator/model/pydantic/__init__.py +++ b/datamodel_code_generator/model/pydantic/__init__.py @@ -19,6 +19,7 @@ class Config(_BaseModel): title: Optional[str] = None allow_population_by_field_name: Optional[bool] = None allow_mutation: Optional[bool] = None + arbitrary_types_allowed: Optional[bool] = None # def get_validator_template() -> Template: diff --git a/datamodel_code_generator/model/pydantic/base_model.py b/datamodel_code_generator/model/pydantic/base_model.py index 0435051d2..bef01de7b 100644 --- a/datamodel_code_generator/model/pydantic/base_model.py +++ b/datamodel_code_generator/model/pydantic/base_model.py @@ -144,6 +144,10 @@ def __init__( config_parameters[config_attribute] = self.extra_template_data[ config_attribute ] + for data_type in self.all_data_types: + if data_type.is_custom_type: + config_parameters['arbitrary_types_allowed'] = True + break if config_parameters: from datamodel_code_generator.model.pydantic import Config diff --git a/datamodel_code_generator/parser/jsonschema.py b/datamodel_code_generator/parser/jsonschema.py index 32c933ad6..a6c30fa64 100644 --- a/datamodel_code_generator/parser/jsonschema.py +++ b/datamodel_code_generator/parser/jsonschema.py @@ -185,6 +185,7 @@ def validate_ref(cls, value: Any) -> Any: examples: Any default: Any id: Optional[str] = Field(default=None, alias='$id') + custom_type_path: Optional[str] = Field(default=None, alias='customTypePath') _raw: Dict[str, Any] class Config: @@ -588,6 +589,10 @@ def parse_item( ) elif item.ref: return self.get_ref_data_type(item.ref) + elif item.custom_type_path: + return self.data_type_manager.get_data_type_from_full_path( + item.custom_type_path, is_custom_type=True + ) elif item.is_array: return self.parse_array_fields( name, item, get_special_path('array', path) @@ -754,6 +759,10 @@ def parse_root_type( ) -> DataType: if obj.ref: data_type: DataType = self.get_ref_data_type(obj.ref) + elif obj.custom_type_path: + data_type = self.data_type_manager.get_data_type_from_full_path( + obj.custom_type_path, is_custom_type=True + ) elif obj.is_object or obj.anyOf or obj.oneOf: data_types: List[DataType] = [] object_path = [*path, name] diff --git a/datamodel_code_generator/types.py b/datamodel_code_generator/types.py index 85cb46d63..800938e7f 100644 --- a/datamodel_code_generator/types.py +++ b/datamodel_code_generator/types.py @@ -70,6 +70,7 @@ class DataType(_BaseModel): is_optional: bool = False is_dict: bool = False is_list: bool = False + is_custom_type: bool = False literals: List[str] = [] use_standard_collections: bool = False use_generic_container: bool = False @@ -89,6 +90,7 @@ def from_import( is_optional: bool = False, is_dict: bool = False, is_list: bool = False, + is_custom_type: bool = False, strict: bool = False, kwargs: Optional[Dict[str, Any]] = None, ) -> 'DataTypeT': @@ -99,6 +101,7 @@ def from_import( is_dict=is_dict, is_list=is_list, is_func=True if kwargs else False, + is_custom_type=is_custom_type, strict=strict, kwargs=kwargs, ) @@ -336,3 +339,10 @@ def __init__( @abstractmethod def get_data_type(self, types: Types, **kwargs: Any) -> DataType: raise NotImplementedError + + def get_data_type_from_full_path( + self, full_path: str, is_custom_type: bool + ) -> DataType: + return self.data_type.from_import( + Import.from_full_path(full_path), is_custom_type=is_custom_type + ) diff --git a/tests/data/expected/main/main_jsonschema_custom_type_path/output.py b/tests/data/expected/main/main_jsonschema_custom_type_path/output.py new file mode 100644 index 000000000..ddead7c6c --- /dev/null +++ b/tests/data/expected/main/main_jsonschema_custom_type_path/output.py @@ -0,0 +1,36 @@ +# generated by datamodel-codegen: +# filename: custom_type_path.json +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from typing import Optional + +from pydantic import BaseModel, Field + +from custom import MultipleLineString, SpecialString, TitleString +from custom.collection.array import Friends +from custom.special import UpperString +from custom.special.numbers import Age + + +class Person(BaseModel): + class Config: + arbitrary_types_allowed = True + + firstName: Optional[TitleString] = Field( + None, description="The person's first name." + ) + lastName: Optional[UpperString] = Field(None, description="The person's last name.") + age: Optional[Age] = Field( + None, description='Age in years which must be equal to or greater than zero.' + ) + friends: Optional[Friends] = None + comment: Optional[MultipleLineString] = None + + +class RootedCustomType(BaseModel): + class Config: + arbitrary_types_allowed = True + + __root__: SpecialString diff --git a/tests/data/jsonschema/custom_type_path.json b/tests/data/jsonschema/custom_type_path.json new file mode 100644 index 000000000..a6d5121bc --- /dev/null +++ b/tests/data/jsonschema/custom_type_path.json @@ -0,0 +1,37 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "Person", + "type": "object", + "properties": { + "firstName": { + "type": "string", + "description": "The person's first name.", + "customTypePath": "custom.TitleString" + }, + "lastName": { + "type": "string", + "description": "The person's last name.", + "customTypePath": "custom.special.UpperString" + }, + "age": { + "description": "Age in years which must be equal to or greater than zero.", + "type": "integer", + "minimum": 0, + "customTypePath": "custom.special.numbers.Age" + }, + "friends": { + "type": "array", + "customTypePath": "custom.collection.array.Friends" + }, + "comment": { + "type": "null", + "customTypePath": "custom.MultipleLineString" + } + }, + "definitions": { + "RootedCustomType": { + "type": "string", + "customTypePath": "custom.SpecialString" + } + } +} \ No newline at end of file diff --git a/tests/test_main.py b/tests/test_main.py index a94209de4..5b87fc6b3 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -4,6 +4,7 @@ from tempfile import TemporaryDirectory from unittest.mock import call +import isort import pytest from _pytest.capture import CaptureFixture from _pytest.tmpdir import TempdirFactory @@ -2884,3 +2885,32 @@ def test_main_jsonschema_field_extras(): ) with pytest.raises(SystemExit): main() + + +@pytest.mark.skipif( + not isort.__version__.startswith('4.'), + reason="isort 5.x don't sort pydantic modules", +) +@freeze_time('2019-07-26') +def test_main_jsonschema_custom_type_path(): + with TemporaryDirectory() as output_dir: + output_file: Path = Path(output_dir) / 'output.py' + return_code: Exit = main( + [ + '--input', + str(JSON_SCHEMA_DATA_PATH / 'custom_type_path.json'), + '--output', + str(output_file), + '--input-file-type', + 'jsonschema', + ] + ) + assert return_code == Exit.OK + assert ( + output_file.read_text() + == ( + EXPECTED_MAIN_PATH / 'main_jsonschema_custom_type_path' / 'output.py' + ).read_text() + ) + with pytest.raises(SystemExit): + main()