Repository: koxudaxi/datamodel-code-generator

Issue:
oneOf schema file $ref not using imports correctly when collapsing root models
**Describe the bug**
When a schema in one file references another schema in another file using the **$ref** keyword, and the referenced schema uses the **oneOf** keyword in the OpenAPI specification, the model generated does not utilize the imports from the referenced schema.

For example. the schema **SupportedCommands** in **commands.yaml** has a $ref to *the *AddUserCommand** schema in **user.yaml.**  The AddUserCommand schema is using the `oneOf` keyword. The generated `commands.py` imports `user` as expected but **doesn't use the import when referencing the user models**

This bug only occurs when the models are generated using an **input directory** of OpenAPI schemas instead of a single openapi.yaml. It also only appears to be happening when I use the `-collapse-root-models` flag. When I generate without that flag, the imports are referenced correctly.

**Generated models with unused imports**
![image](https://github.com/koxudaxi/datamodel-code-generator/assets/21082857/52f4db90-e1ac-4e89-933e-072c6a5dc553)

**To Reproduce**

1. Create input directory called schemas, containing two files, `user.yaml` and `commands.yaml` using the schemas below
2. Create output directory called models
3. Generate models using the datamodel-codegen command below in the root directory of the project

Example  **user.yaml** schema:

```yaml
openapi: 3.0.0
components:
  schemas:
    JsonRpcError:
      type: object
      properties:
        error:
          type: string
    User:
      type: object
      properties:
        userId:
          type: string
    AddUserCommand:
      oneOf:
        - $ref: "#/components/schemas/User"
        - $ref: "#/components/schemas/JsonRpcError"
```

Example  **commands.yaml** schema:
```yaml
openapi: 3.0.0
components:
  schemas:
    SupportedCommands:
      $ref: "user.yaml#/components/schemas/AddUserCommand"
```

Used commandline:
```
$ datamodel-codegen --input schemas --output models --collapse-root-models
```

**Expected behavior**
The user models referenced in commands.py should be correctly prefixed with the imported module name. For example, instead of `User` / `JsonRpcError`, it should be `user.User` / `user.JsonRpcError`,  like shown below:

```python3
# generated by datamodel-codegen:
#   filename:  commands.yaml
#   timestamp: 2023-07-31T22:40:48+00:00

from __future__ import annotations

from pydantic import BaseModel

from . import user


class SupportedCommands(BaseModel):
    __root__: Union[user.User, user.JsonRpcError] 
```


**Version:**
- OS: Ubuntu 22.04.1 LTS
- Python version: Python 3.10.6
- datamodel-code-generator version: 0.21.2




Discussion:

Patch:
diff --git a/src/datamodel_code_generator/parser/base.py b/src/datamodel_code_generator/parser/base.py
index 32c771f8d..0bb947ef3 100644
--- a/src/datamodel_code_generator/parser/base.py
+++ b/src/datamodel_code_generator/parser/base.py
@@ -935,6 +935,8 @@ def __collapse_root_models(  # noqa: PLR0912
                 for data_type in model_field.data_type.all_data_types:
                     reference = data_type.reference
                     if not reference or not isinstance(reference.source, self.data_model_root_type):
+                        # If the data type is not a reference, we can't collapse it.
+                        # If it's a reference to a root model type, we don't do anything.
                         continue
 
                     # Use root-type as model_field type
@@ -949,6 +951,10 @@ def __collapse_root_models(  # noqa: PLR0912
                     ):
                         continue  # pragma: no cover
 
+                    if root_type_field.data_type.reference:
+                        # If the root type field is a reference, we aren't able to collapse it yet.
+                        continue
+
                     # set copied data_type
                     copied_data_type = root_type_field.data_type.copy()
                     if isinstance(data_type.parent, self.data_model_field_type):


Test patch:
diff --git a/tests/data/expected/main/jsonschema/forwarding_reference/__init__.py b/tests/data/expected/main/jsonschema/forwarding_reference/__init__.py
new file mode 100644
index 000000000..7cf901d0c
--- /dev/null
+++ b/tests/data/expected/main/jsonschema/forwarding_reference/__init__.py
@@ -0,0 +1,3 @@
+# generated by datamodel-codegen:
+#   filename:  forwarding_reference
+#   timestamp: 2019-07-26T00:00:00+00:00
diff --git a/tests/data/expected/main/jsonschema/forwarding_reference/commons.py b/tests/data/expected/main/jsonschema/forwarding_reference/commons.py
new file mode 100644
index 000000000..737a97037
--- /dev/null
+++ b/tests/data/expected/main/jsonschema/forwarding_reference/commons.py
@@ -0,0 +1,13 @@
+# generated by datamodel-codegen:
+#   filename:  commons.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class Model(BaseModel):
+    __root__: Any
diff --git a/tests/data/expected/main/jsonschema/forwarding_reference/forwarding.py b/tests/data/expected/main/jsonschema/forwarding_reference/forwarding.py
new file mode 100644
index 000000000..dba314a02
--- /dev/null
+++ b/tests/data/expected/main/jsonschema/forwarding_reference/forwarding.py
@@ -0,0 +1,17 @@
+# generated by datamodel-codegen:
+#   filename:  forwarding.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Any, List
+
+from pydantic import BaseModel
+
+
+class Model(BaseModel):
+    __root__: Any
+
+
+class ForwardingArray(BaseModel):
+    __root__: List
diff --git a/tests/data/expected/main/jsonschema/forwarding_reference/schema.py b/tests/data/expected/main/jsonschema/forwarding_reference/schema.py
new file mode 100644
index 000000000..32d2b7e00
--- /dev/null
+++ b/tests/data/expected/main/jsonschema/forwarding_reference/schema.py
@@ -0,0 +1,19 @@
+# generated by datamodel-codegen:
+#   filename:  schema.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel
+
+from . import forwarding
+
+
+class Model(BaseModel):
+    __root__: Any
+
+
+class Pets(BaseModel):
+    __root__: forwarding.ForwardingArray
diff --git a/tests/data/jsonschema/forwarding_reference/commons.json b/tests/data/jsonschema/forwarding_reference/commons.json
new file mode 100644
index 000000000..4d27bcf05
--- /dev/null
+++ b/tests/data/jsonschema/forwarding_reference/commons.json
@@ -0,0 +1,6 @@
+{
+  "$defs": {
+    "customArray": { "type": "array" }
+  },
+  "$schema": "https://json-schema.org/draft/2020-12/schema"
+}
diff --git a/tests/data/jsonschema/forwarding_reference/forwarding.json b/tests/data/jsonschema/forwarding_reference/forwarding.json
new file mode 100644
index 000000000..28b4ce11e
--- /dev/null
+++ b/tests/data/jsonschema/forwarding_reference/forwarding.json
@@ -0,0 +1,6 @@
+{
+  "$defs": {
+    "forwardingArray": {"$ref": "commons.json#/$defs/customArray"}
+  },
+  "$schema": "https://json-schema.org/draft/2020-12/schema"
+}
diff --git a/tests/data/jsonschema/forwarding_reference/schema.json b/tests/data/jsonschema/forwarding_reference/schema.json
new file mode 100644
index 000000000..bff8245d8
--- /dev/null
+++ b/tests/data/jsonschema/forwarding_reference/schema.json
@@ -0,0 +1,6 @@
+{
+  "$defs": {
+    "pets": {"$ref": "forwarding.json#/$defs/forwardingArray"}
+  },
+  "$schema": "https://json-schema.org/draft/2020-12/schema"
+}
diff --git a/tests/main/jsonschema/test_main_jsonschema.py b/tests/main/jsonschema/test_main_jsonschema.py
index f4309ce25..41491da14 100644
--- a/tests/main/jsonschema/test_main_jsonschema.py
+++ b/tests/main/jsonschema/test_main_jsonschema.py
@@ -3143,3 +3143,24 @@ def test_main_jsonschema_same_name_objects(tmp_path: Path) -> None:
     ])
     assert return_code == Exit.OK
     assert output_file.read_text() == (EXPECTED_JSON_SCHEMA_PATH / "same_name_objects.py").read_text()
+
+
+@freeze_time("2019-07-26")
+def test_main_jsonschema_forwarding_reference_collapse_root(tmp_path: Path) -> None:
+    """
+    See: https://github.com/koxudaxi/datamodel-code-generator/issues/1466
+    """
+    return_code: Exit = main([
+        "--input",
+        str(JSON_SCHEMA_DATA_PATH / "forwarding_reference"),
+        "--output",
+        str(tmp_path),
+        "--input-file-type",
+        "jsonschema",
+        "--collapse-root-models",
+    ])
+    assert return_code == Exit.OK
+    main_modular_dir = EXPECTED_JSON_SCHEMA_PATH / "forwarding_reference"
+    for path in main_modular_dir.rglob("*.py"):
+        result = tmp_path.joinpath(path.relative_to(main_modular_dir)).read_text()
+        assert result == path.read_text()