From eb5e4435664382d2ea1e1f8f0bdd731aaeda78d3 Mon Sep 17 00:00:00 2001 From: Kevin Michel Date: Tue, 23 Jul 2024 14:23:31 +0200 Subject: [PATCH] Fix missing imports for collapsed models When an externally referenced model is collapsed, the required imports for the replacing type(s) were missing. --- datamodel_code_generator/parser/base.py | 36 +++++++++++++++++-- .../expected/main/jsonschema/external_ref0.py | 15 ++++++++ .../jsonschema/external_reference/ref0.json | 8 +++++ .../jsonschema/external_reference/ref1.json | 9 +++++ .../jsonschema/external_reference/ref2.json | 6 ++++ tests/main/jsonschema/test_main_jsonschema.py | 23 ++++++++++++ 6 files changed, 95 insertions(+), 2 deletions(-) create mode 100644 tests/data/expected/main/jsonschema/external_ref0.py create mode 100644 tests/data/jsonschema/external_reference/ref0.json create mode 100644 tests/data/jsonschema/external_reference/ref1.json create mode 100644 tests/data/jsonschema/external_reference/ref2.json diff --git a/datamodel_code_generator/parser/base.py b/datamodel_code_generator/parser/base.py index f359adbf0..4d7ddedc1 100644 --- a/datamodel_code_generator/parser/base.py +++ b/datamodel_code_generator/parser/base.py @@ -958,7 +958,11 @@ def __reuse_model( models.remove(duplicate) def __collapse_root_models( - self, models: List[DataModel], unused_models: List[DataModel], imports: Imports + self, + models: List[DataModel], + unused_models: List[DataModel], + imports: Imports, + scoped_model_resolver: ModelResolver, ) -> None: if not self.collapse_root_models: return None @@ -1033,6 +1037,32 @@ def __collapse_root_models( ] else: # pragma: no cover continue + + for d in root_type_field.data_type.data_types: + if d.reference is None: + continue + from_, import_ = full_path = relative( + model.module_name, d.full_name + ) + if from_ and import_: + alias = scoped_model_resolver.add(full_path, import_).name + name = d.reference.short_name + if alias != name: + d.alias = ( + alias + if d.reference.short_name == import_ + else f'{alias}.{name}' + ) + imports.append( + [ + Import( + from_=from_, + import_=import_, + reference_path=d.reference.path, + ) + ] + ) + original_field = get_most_of_parent(data_type, DataModelFieldBase) if original_field: # pragma: no cover # TODO: Improve detection of reference type @@ -1329,7 +1359,9 @@ class Processed(NamedTuple): self.__extract_inherited_enum(models) self.__set_reference_default_value_to_field(models) self.__reuse_model(models, require_update_action_models) - self.__collapse_root_models(models, unused_models, imports) + self.__collapse_root_models( + models, unused_models, imports, scoped_model_resolver + ) self.__set_default_enum_member(models) self.__sort_models(models, imports) self.__set_one_literal_on_default(models) diff --git a/tests/data/expected/main/jsonschema/external_ref0.py b/tests/data/expected/main/jsonschema/external_ref0.py new file mode 100644 index 000000000..f6f3ad63e --- /dev/null +++ b/tests/data/expected/main/jsonschema/external_ref0.py @@ -0,0 +1,15 @@ +# generated by datamodel-codegen: +# filename: ref0.json +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from typing import Optional + +from pydantic import BaseModel + +from . import ref2 + + +class Model(BaseModel): + ref1: Optional[ref2.Model] = None diff --git a/tests/data/jsonschema/external_reference/ref0.json b/tests/data/jsonschema/external_reference/ref0.json new file mode 100644 index 000000000..72b7d7ac1 --- /dev/null +++ b/tests/data/jsonschema/external_reference/ref0.json @@ -0,0 +1,8 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "properties": { + "ref1": { + "$ref": "ref1.json#/" + } + } +} diff --git a/tests/data/jsonschema/external_reference/ref1.json b/tests/data/jsonschema/external_reference/ref1.json new file mode 100644 index 000000000..79f70e4c9 --- /dev/null +++ b/tests/data/jsonschema/external_reference/ref1.json @@ -0,0 +1,9 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "anyOf": [ + { + "$ref": "ref2.json#/" + }, + {"type": "null"} + ] +} diff --git a/tests/data/jsonschema/external_reference/ref2.json b/tests/data/jsonschema/external_reference/ref2.json new file mode 100644 index 000000000..46e6ea860 --- /dev/null +++ b/tests/data/jsonschema/external_reference/ref2.json @@ -0,0 +1,6 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "properties": { + "key": {"type": "string"} + } +} diff --git a/tests/main/jsonschema/test_main_jsonschema.py b/tests/main/jsonschema/test_main_jsonschema.py index 6409869d5..c2e3c8f71 100644 --- a/tests/main/jsonschema/test_main_jsonschema.py +++ b/tests/main/jsonschema/test_main_jsonschema.py @@ -244,6 +244,29 @@ def test_main_jsonschema_external_files(): ) +@pytest.mark.benchmark +@freeze_time('2019-07-26') +def test_main_jsonschema_collapsed_external_references(): + with TemporaryDirectory() as output_dir: + output_dir: Path = Path(output_dir) / 'output' + output_dir.mkdir() + return_code: Exit = main( + [ + '--input', + str(JSON_SCHEMA_DATA_PATH / 'external_reference'), + '--output', + str(output_dir), + '--input-file-type', + 'jsonschema', + '--collapse-root-models', + ] + ) + assert return_code == Exit.OK + assert (output_dir / 'ref0.py').read_text() == ( + EXPECTED_JSON_SCHEMA_PATH / 'external_ref0.py' + ).read_text() + + @pytest.mark.benchmark @freeze_time('2019-07-26') def test_main_jsonschema_multiple_files():