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

@@ -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)

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

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#/"
+        }
+    }
+}

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"}
+    ]
+}

tests/data/jsonschema/external_reference/ref2.json

@@ -0,0 +1,6 @@
+{
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "properties": {
+        "key": {"type": "string"}
+    }
+}

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():