diff --git a/datamodel_code_generator/imports.py b/datamodel_code_generator/imports.py index 981c1b15e..2d5f79121 100644 --- a/datamodel_code_generator/imports.py +++ b/datamodel_code_generator/imports.py @@ -47,3 +47,5 @@ def append(self, imports: Union[Import, List[Import], None]) -> None: IMPORT_ENUM = Import(import_='Enum', from_='enum') IMPORT_ANNOTATIONS = Import(from_='__future__', import_='annotations') IMPORT_CONSTR = Import(import_='constr', from_='pydantic') +IMPORT_CONINT = Import(import_='conint', from_='pydantic') +IMPORT_CONFLOAT = Import(import_='confloat', from_='pydantic') diff --git a/datamodel_code_generator/model/pydantic/base_model.py b/datamodel_code_generator/model/pydantic/base_model.py index 70756e94a..08fa85c3b 100644 --- a/datamodel_code_generator/model/pydantic/base_model.py +++ b/datamodel_code_generator/model/pydantic/base_model.py @@ -27,6 +27,7 @@ def __init__( extra_template_data: Optional[DefaultDict[str, Any]] = None, auto_import: bool = True, reference_classes: Optional[List[str]] = None, + imports: Optional[List[Import]] = None, ): super().__init__( @@ -39,6 +40,7 @@ def __init__( extra_template_data=extra_template_data, auto_import=auto_import, reference_classes=reference_classes, + imports=imports, ) if 'additionalProperties' in self.extra_template_data: diff --git a/datamodel_code_generator/model/pydantic/types.py b/datamodel_code_generator/model/pydantic/types.py index a47d72945..a5e82f05f 100644 --- a/datamodel_code_generator/model/pydantic/types.py +++ b/datamodel_code_generator/model/pydantic/types.py @@ -1,6 +1,11 @@ from typing import Any, Dict -from datamodel_code_generator.imports import IMPORT_CONSTR, Import +from datamodel_code_generator.imports import ( + IMPORT_CONFLOAT, + IMPORT_CONINT, + IMPORT_CONSTR, + Import, +) from datamodel_code_generator.types import DataType, Types type_map: Dict[Types, DataType] = { @@ -80,7 +85,12 @@ def get_data_int_type(types: Types, **kwargs: Any) -> DataType: return DataType(type='PositiveInt') if len(data_type_kwargs) == 1 and data_type_kwargs.get('ge') == 0: return DataType(type='NegativeInt') - return DataType(type='conint', is_func=True, kwargs=data_type_kwargs) + return DataType( + type='conint', + is_func=True, + kwargs=data_type_kwargs, + imports_=[IMPORT_CONINT], + ) return type_map[types] @@ -102,7 +112,12 @@ def get_data_float_type(types: Types, **kwargs: Any) -> DataType: return DataType(type='PositiveFloat') if len(data_type_kwargs) == 1 and data_type_kwargs.get('ge') == 0: return DataType(type='NegativeFloat') - return DataType(type='confloat', is_func=True, kwargs=data_type_kwargs) + return DataType( + type='confloat', + is_func=True, + kwargs=data_type_kwargs, + imports_=[IMPORT_CONFLOAT], + ) return type_map[types] diff --git a/datamodel_code_generator/parser/openapi.py b/datamodel_code_generator/parser/openapi.py index 7a1b40008..8cb8c68c8 100644 --- a/datamodel_code_generator/parser/openapi.py +++ b/datamodel_code_generator/parser/openapi.py @@ -144,6 +144,9 @@ def parse_all_of(self, name: str, obj: JsonSchemaObject) -> List[DataType]: custom_template_dir=self.custom_template_dir, extra_template_data=self.extra_template_data, ) + # add imports for the fields + for f in fields: + data_model_type.imports.extend(f.imports) self.append_result(data_model_type) return [self.data_type(type=name, ref=True, version_compatible=True)] diff --git a/tests/data/allof.yaml b/tests/data/allof.yaml index a0a5da289..7f1064602 100644 --- a/tests/data/allof.yaml +++ b/tests/data/allof.yaml @@ -131,6 +131,17 @@ components: properties: number: type: string + AllOfCombine: + allOf: + - $ref: "#/components/schemas/Pet" + - type: object + properties: + birthdate: + type: string + format: date + size: + type: integer + minimum: 1 AnyOfCombine: allOf: - $ref: "#/components/schemas/Pet" @@ -168,6 +179,9 @@ components: properties: age: type: string + birthdate: + type: string + format: date-time Error: required: - code diff --git a/tests/data/anyof.yaml b/tests/data/anyof.yaml index 5f3c03a5b..eebc449dd 100644 --- a/tests/data/anyof.yaml +++ b/tests/data/anyof.yaml @@ -153,6 +153,9 @@ components: properties: name: type: string + birthday: + type: string + format: date Error: required: - code diff --git a/tests/model/pydantic/test_types.py b/tests/model/pydantic/test_types.py index 23880f0a5..6824ed755 100644 --- a/tests/model/pydantic/test_types.py +++ b/tests/model/pydantic/test_types.py @@ -1,5 +1,9 @@ import pytest -from datamodel_code_generator.imports import IMPORT_CONSTR +from datamodel_code_generator.imports import ( + IMPORT_CONFLOAT, + IMPORT_CONINT, + IMPORT_CONSTR, +) from datamodel_code_generator.model.pydantic.types import ( get_data_float_type, get_data_int_type, @@ -16,27 +20,40 @@ ( Types.integer, {'maximum': 10}, - DataType(type='conint', is_func=True, kwargs={'gt': 10}), + DataType( + type='conint', is_func=True, kwargs={'gt': 10}, imports_=[IMPORT_CONINT] + ), ), ( Types.integer, {'exclusiveMaximum': 10}, - DataType(type='conint', is_func=True, kwargs={'ge': 10}), + DataType( + type='conint', is_func=True, kwargs={'ge': 10}, imports_=[IMPORT_CONINT] + ), ), ( Types.integer, {'minimum': 10}, - DataType(type='conint', is_func=True, kwargs={'lt': 10}), + DataType( + type='conint', is_func=True, kwargs={'lt': 10}, imports_=[IMPORT_CONINT] + ), ), ( Types.integer, {'exclusiveMinimum': 10}, - DataType(type='conint', is_func=True, kwargs={'le': 10}), + DataType( + type='conint', is_func=True, kwargs={'le': 10}, imports_=[IMPORT_CONINT] + ), ), ( Types.integer, {'multipleOf': 10}, - DataType(type='conint', is_func=True, kwargs={'multiple_of': 10}), + DataType( + type='conint', + is_func=True, + kwargs={'multiple_of': 10}, + imports_=[IMPORT_CONINT], + ), ), (Types.integer, {'exclusiveMinimum': 0}, DataType(type='PositiveInt')), (Types.integer, {'exclusiveMaximum': 0}, DataType(type='NegativeInt')), @@ -53,27 +70,52 @@ def test_get_data_int_type(types, params, data_type): ( Types.float, {'maximum': 10}, - DataType(type='confloat', is_func=True, kwargs={'gt': 10}), + DataType( + type='confloat', + is_func=True, + kwargs={'gt': 10}, + imports_=[IMPORT_CONFLOAT], + ), ), ( Types.float, {'exclusiveMaximum': 10}, - DataType(type='confloat', is_func=True, kwargs={'ge': 10}), + DataType( + type='confloat', + is_func=True, + kwargs={'ge': 10}, + imports_=[IMPORT_CONFLOAT], + ), ), ( Types.float, {'minimum': 10}, - DataType(type='confloat', is_func=True, kwargs={'lt': 10}), + DataType( + type='confloat', + is_func=True, + kwargs={'lt': 10}, + imports_=[IMPORT_CONFLOAT], + ), ), ( Types.float, {'exclusiveMinimum': 10}, - DataType(type='confloat', is_func=True, kwargs={'le': 10}), + DataType( + type='confloat', + is_func=True, + kwargs={'le': 10}, + imports_=[IMPORT_CONFLOAT], + ), ), ( Types.float, {'multipleOf': 10}, - DataType(type='confloat', is_func=True, kwargs={'multiple_of': 10}), + DataType( + type='confloat', + is_func=True, + kwargs={'multiple_of': 10}, + imports_=[IMPORT_CONFLOAT], + ), ), (Types.float, {'exclusiveMinimum': 0}, DataType(type='PositiveFloat')), (Types.float, {'exclusiveMaximum': 0}, DataType(type='NegativeFloat')), diff --git a/tests/parser/test_openapi.py b/tests/parser/test_openapi.py index b078b70a4..6657cd1a5 100644 --- a/tests/parser/test_openapi.py +++ b/tests/parser/test_openapi.py @@ -771,6 +771,7 @@ def test_openapi_parser_parse_anyof(): parser.parse() == '''from __future__ import annotations +from datetime import date from typing import List, Optional, Union from pydantic import BaseModel @@ -806,6 +807,7 @@ class AnyOfobj(BaseModel): class AnyOfArrayItem(BaseModel): name: Optional[str] = None + birthday: Optional[date] = None class AnyOfArray(BaseModel): @@ -827,9 +829,10 @@ def test_openapi_parser_parse_allof(): parser.parse() == '''from __future__ import annotations +from datetime import date, datetime from typing import List, Optional -from pydantic import BaseModel +from pydantic import BaseModel, conint class Pet(BaseModel): @@ -851,6 +854,11 @@ class AllOfobj(BaseModel): number: Optional[str] = None +class AllOfCombine(Pet): + birthdate: Optional[date] = None + size: Optional[conint(lt=1.0)] = None + + class AnyOfCombine(Pet, Car): age: Optional[str] = None @@ -873,6 +881,7 @@ class AnyOfCombineInArray(BaseModel): class AnyOfCombineInRoot(Pet, Car): age: Optional[str] = None + birthdate: Optional[datetime] = None class Error(BaseModel): diff --git a/tests/test_main_kr.py b/tests/test_main_kr.py new file mode 100644 index 000000000..d6ef89cf9 --- /dev/null +++ b/tests/test_main_kr.py @@ -0,0 +1,717 @@ +import shutil +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Mapping + +import pytest +from _pytest.capture import CaptureFixture +from _pytest.tmpdir import TempdirFactory +from datamodel_code_generator.__main__ import Exit, main +from freezegun import freeze_time + +DATA_PATH: Path = Path(__file__).parent / 'data' +TIMESTAMP = '1985-10-26T01:21:00-07:00' + + +@freeze_time('2019-07-26') +def test_main(): + with TemporaryDirectory() as output_dir: + output_file: Path = Path(output_dir) / 'output.py' + return_code: Exit = main( + ['--input', str(DATA_PATH / 'api.yaml'), '--output', str(output_file)] + ) + assert return_code == Exit.OK + assert ( + output_file.read_text() + == '''# generated by datamodel-codegen: +# filename: api.yaml +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import BaseModel, UrlStr + + +class Pet(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Pets(BaseModel): + __root__: List[Pet] + + +class User(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Users(BaseModel): + __root__: List[User] + + +class Id(BaseModel): + __root__: str + + +class Rules(BaseModel): + __root__: List[str] + + +class Error(BaseModel): + code: int + message: str + + +class api(BaseModel): + apiKey: Optional[str] = None + apiVersionNumber: Optional[str] = None + apiUrl: Optional[UrlStr] = None + apiDocumentationUrl: Optional[UrlStr] = None + + +class apis(BaseModel): + __root__: List[api] + + +class Event(BaseModel): + name: Optional[str] = None + + +class Result(BaseModel): + event: Optional[Event] = None +''' + ) + + with pytest.raises(SystemExit): + main() + + +@freeze_time('2019-07-26') +def test_main_base_class(): + with TemporaryDirectory() as output_dir: + output_file: Path = Path(output_dir) / 'output.py' + return_code: Exit = main( + [ + '--input', + str(DATA_PATH / 'api.yaml'), + '--output', + str(output_file), + '--base-class', + 'custom_module.Base', + ] + ) + assert return_code == Exit.OK + assert ( + output_file.read_text() + == '''# generated by datamodel-codegen: +# filename: api.yaml +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import UrlStr + +from custom_module import Base + + +class Pet(Base): + id: int + name: str + tag: Optional[str] = None + + +class Pets(Base): + __root__: List[Pet] + + +class User(Base): + id: int + name: str + tag: Optional[str] = None + + +class Users(Base): + __root__: List[User] + + +class Id(Base): + __root__: str + + +class Rules(Base): + __root__: List[str] + + +class Error(Base): + code: int + message: str + + +class api(Base): + apiKey: Optional[str] = None + apiVersionNumber: Optional[str] = None + apiUrl: Optional[UrlStr] = None + apiDocumentationUrl: Optional[UrlStr] = None + + +class apis(Base): + __root__: List[api] + + +class Event(Base): + name: Optional[str] = None + + +class Result(Base): + event: Optional[Event] = None +''' + ) + + with pytest.raises(SystemExit): + main() + + +@freeze_time('2019-07-26') +def test_target_python_version(): + with TemporaryDirectory() as output_dir: + output_file: Path = Path(output_dir) / 'output.py' + return_code: Exit = main( + [ + '--input', + str(DATA_PATH / 'api.yaml'), + '--output', + str(output_file), + '--target-python-version', + '3.6', + ] + ) + assert return_code == Exit.OK + assert ( + output_file.read_text() + == '''# generated by datamodel-codegen: +# filename: api.yaml +# timestamp: 2019-07-26T00:00:00+00:00 + +from typing import List, Optional + +from pydantic import BaseModel, UrlStr + + +class Pet(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Pets(BaseModel): + __root__: List['Pet'] + + +class User(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Users(BaseModel): + __root__: List['User'] + + +class Id(BaseModel): + __root__: str + + +class Rules(BaseModel): + __root__: List[str] + + +class Error(BaseModel): + code: int + message: str + + +class api(BaseModel): + apiKey: Optional[str] = None + apiVersionNumber: Optional[str] = None + apiUrl: Optional[UrlStr] = None + apiDocumentationUrl: Optional[UrlStr] = None + + +class apis(BaseModel): + __root__: List['api'] + + +class Event(BaseModel): + name: Optional[str] = None + + +class Result(BaseModel): + event: Optional['Event'] = None +''' + ) + + with pytest.raises(SystemExit): + main() + + +@pytest.mark.parametrize( + 'expected', + [ + { + ( + '__init__.py', + ): '''\ +# generated by datamodel-codegen: +# filename: modular.yaml +# timestamp: 1985-10-26T08:21:00+00:00 + +from __future__ import annotations + +from typing import Optional + +from pydantic import BaseModel + +from . import models + + +class Id(BaseModel): + __root__: str + + +class Error(BaseModel): + code: int + message: str + + +class Result(BaseModel): + event: Optional[models.Event] = None +''', + ( + 'models.py', + ): '''\ +# generated by datamodel-codegen: +# filename: modular.yaml +# timestamp: 1985-10-26T08:21:00+00:00 + +from __future__ import annotations + +from enum import Enum +from typing import Any, Dict, Optional, Union + +from pydantic import BaseModel + + +class Species(Enum): + dog = 'dog' + cat = 'cat' + snake = 'snake' + + +class Pet(BaseModel): + id: int + name: str + tag: Optional[str] = None + species: Optional[Species] = None + + +class User(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Event(BaseModel): + name: Optional[Union[str, float, int, bool, Dict[str, Any]]] = None +''', + ( + 'collections.py', + ): '''\ +# generated by datamodel-codegen: +# filename: modular.yaml +# timestamp: 1985-10-26T08:21:00+00:00 + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import BaseModel, UrlStr + +from . import models + + +class Pets(BaseModel): + __root__: List[models.Pet] + + +class Users(BaseModel): + __root__: List[models.User] + + +class Rules(BaseModel): + __root__: List[str] + + +class api(BaseModel): + apiKey: Optional[str] = None + apiVersionNumber: Optional[str] = None + apiUrl: Optional[UrlStr] = None + apiDocumentationUrl: Optional[UrlStr] = None + + +class apis(BaseModel): + __root__: List[api] +''', + ( + 'foo', + '__init__.py', + ): '''\ +# generated by datamodel-codegen: +# filename: modular.yaml +# timestamp: 1985-10-26T08:21:00+00:00 +''', + ( + 'foo', + 'bar.py', + ): '''\ +# generated by datamodel-codegen: +# filename: modular.yaml +# timestamp: 1985-10-26T08:21:00+00:00 + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel + + +class Thing(BaseModel): + attributes: Optional[Dict[str, Any]] = None + + +class Thang(BaseModel): + attributes: Optional[List[Dict[str, Any]]] = None + + +class Clone(Thing): + pass +''', + } + ], +) +def test_main_modular( + tmpdir_factory: TempdirFactory, expected: Mapping[str, str] +) -> None: + """Test main function on modular file.""" + + output_directory = Path(tmpdir_factory.mktemp('output')) + + input_filename = DATA_PATH / 'modular.yaml' + output_path = output_directory / 'model' + + with freeze_time(TIMESTAMP): + main(['--input', str(input_filename), '--output', str(output_path)]) + + for key, value in expected.items(): + result = output_path.joinpath(*key).read_text() + assert result == value + + +def test_main_modular_no_file() -> None: + """Test main function on modular file with no output name.""" + + input_filename = DATA_PATH / 'modular.yaml' + + assert main(['--input', str(input_filename)]) == Exit.ERROR + + +def test_main_modular_filename(tmpdir_factory: TempdirFactory) -> None: + """Test main function on modular file with filename.""" + + output_directory = Path(tmpdir_factory.mktemp('output')) + + input_filename = DATA_PATH / 'modular.yaml' + output_filename = output_directory / 'model.py' + + assert ( + main(['--input', str(input_filename), '--output', str(output_filename)]) + == Exit.ERROR + ) + + +@pytest.mark.parametrize( + 'expected', + [ + '''\ +# generated by datamodel-codegen: +# filename: api.yaml +# timestamp: 1985-10-26T08:21:00+00:00 + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import BaseModel, UrlStr + + +class Pet(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Pets(BaseModel): + __root__: List[Pet] + + +class User(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Users(BaseModel): + __root__: List[User] + + +class Id(BaseModel): + __root__: str + + +class Rules(BaseModel): + __root__: List[str] + + +class Error(BaseModel): + code: int + message: str + + +class api(BaseModel): + apiKey: Optional[str] = None + apiVersionNumber: Optional[str] = None + apiUrl: Optional[UrlStr] = None + apiDocumentationUrl: Optional[UrlStr] = None + + +class apis(BaseModel): + __root__: List[api] + + +class Event(BaseModel): + name: Optional[str] = None + + +class Result(BaseModel): + event: Optional[Event] = None +''' + ], +) +def test_main_no_file(capsys: CaptureFixture, expected: str) -> None: + """Test main function on non-modular file with no output name.""" + + input_filename = DATA_PATH / 'api.yaml' + + with freeze_time(TIMESTAMP): + main(['--input', str(input_filename)]) + + captured = capsys.readouterr() + assert captured.out == expected + assert not captured.err + + +@pytest.mark.parametrize( + 'expected', + [ + '''\ +# generated by datamodel-codegen: +# filename: api.yaml +# timestamp: 1985-10-26T08:21:00+00:00 + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import BaseModel, UrlStr + + +class Pet(BaseModel): # 1 2, 1 2, this is just a pet + id: int + name: str + tag: Optional[str] = None + + +class Pets(BaseModel): + __root__: List[Pet] + + +class User(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Users(BaseModel): + __root__: List[User] + + +class Id(BaseModel): + __root__: str + + +class Rules(BaseModel): + __root__: List[str] + + +class Error(BaseModel): + code: int + message: str + + +class api(BaseModel): + apiKey: Optional[str] = None + apiVersionNumber: Optional[str] = None + apiUrl: Optional[UrlStr] = None + apiDocumentationUrl: Optional[UrlStr] = None + + +class apis(BaseModel): + __root__: List[api] + + +class Event(BaseModel): + name: Optional[str] = None + + +class Result(BaseModel): + event: Optional[Event] = None +''' + ], +) +def test_main_custom_template_dir(capsys: CaptureFixture, expected: str) -> None: + """Test main function with custom template directory.""" + + input_filename = DATA_PATH / 'api.yaml' + custom_template_dir = DATA_PATH / 'templates' + extra_template_data = DATA_PATH / 'extra_data.json' + + with freeze_time(TIMESTAMP): + main( + [ + '--input', + str(input_filename), + '--custom-template-dir', + str(custom_template_dir), + '--extra-template-data', + str(extra_template_data), + ] + ) + + captured = capsys.readouterr() + assert captured.out == expected + assert not captured.err + + +@freeze_time('2019-07-26') +def test_pyproject(): + with TemporaryDirectory() as output_dir: + output_dir = Path(output_dir) + pyproject_toml = Path(DATA_PATH) / "project" / "pyproject.toml" + shutil.copy(pyproject_toml, output_dir) + output_file: Path = output_dir / 'output.py' + return_code: Exit = main( + ['--input', str(DATA_PATH / 'api.yaml'), '--output', str(output_file)] + ) + assert return_code == Exit.OK + assert ( + output_file.read_text() + == '''# generated by datamodel-codegen: +# filename: api.yaml +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import ( + annotations, +) + +from typing import ( + List, + Optional, +) + +from pydantic import ( + BaseModel, + UrlStr, +) + + +class Pet(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Pets(BaseModel): + __root__: List[Pet] + + +class User(BaseModel): + id: int + name: str + tag: Optional[str] = None + + +class Users(BaseModel): + __root__: List[User] + + +class Id(BaseModel): + __root__: str + + +class Rules(BaseModel): + __root__: List[str] + + +class Error(BaseModel): + code: int + message: str + + +class api(BaseModel): + apiKey: Optional[ + str + ] = None + apiVersionNumber: Optional[ + str + ] = None + apiUrl: Optional[ + UrlStr + ] = None + apiDocumentationUrl: Optional[ + UrlStr + ] = None + + +class apis(BaseModel): + __root__: List[api] + + +class Event(BaseModel): + name: Optional[str] = None + + +class Result(BaseModel): + event: Optional[ + Event + ] = None +''' + ) + + with pytest.raises(SystemExit): + main()