Fix: Resolve invalid imports in single file output by improving reference checks (#2604)

Author: koxudaxi

src/datamodel_code_generator/parser/base.py

@@ -792,6 +792,8 @@ def __change_from_import(
         scoped_model_resolver: ModelResolver,
         init: bool,  # noqa: FBT001
     ) -> None:
+        model_paths = {model.path for model in models}
+
         for model in models:
             scoped_model_resolver.add([model.path], model.class_name)
         for model in models:
@@ -800,7 +802,7 @@ def __change_from_import(
             for data_type in model.all_data_types:
                 # To change from/import
 
-                if not data_type.reference or data_type.reference.source in models:
+                if not data_type.reference or data_type.reference.path in model_paths:
                     # No need to import non-reference model.
                     # Or, Referenced model is in the same file. we don't need to import the model
                     continue

tests/data/expected/main/jsonschema/all_of_any_of_base_class_ref.py

@@ -0,0 +1,74 @@
+# generated by datamodel-codegen:
+#   filename:  all_of_any_of_base_class_ref.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional, Union
+
+from pydantic import BaseModel, Field, confloat
+
+
+class MapState1(BaseModel):
+    map_view_mode: str = Field("MODE_2D", alias="mapViewMode", const=True)
+
+
+class MapState2(BaseModel):
+    latitude: Latitude
+    longitude: Longitude
+    zoom: Optional[Zoom] = 0
+    bearing: Optional[Bearing] = None
+    pitch: Pitch
+    drag_rotate: Optional[DragRotate] = Field(None, alias="dragRotate")
+    map_split_mode: str = Field("SWIPE_COMPARE", alias="mapSplitMode", const=True)
+    is_split: bool = Field(True, alias="isSplit", const=True)
+
+
+class MapState3(BaseModel):
+    pass
+
+
+class MapState4(MapState1, MapState3):
+    pass
+
+
+class MapState5(MapState2, MapState3):
+    pass
+
+
+class MapState6(MapState4):
+    pass
+
+
+class MapState7(MapState5):
+    pass
+
+
+class MapState(BaseModel):
+    __root__: Union[MapState4, MapState5, MapState6, MapState7] = Field(
+        ..., title="MapState"
+    )
+
+
+class Bearing(BaseModel):
+    __root__: float
+
+
+class DragRotate(BaseModel):
+    __root__: bool
+
+
+class Latitude(BaseModel):
+    __root__: confloat(ge=-90.0, le=90.0)
+
+
+class Longitude(BaseModel):
+    __root__: confloat(ge=-180.0, le=180.0)
+
+
+class Pitch(BaseModel):
+    __root__: confloat(ge=0.0, lt=90.0)
+
+
+class Zoom(BaseModel):
+    __root__: confloat(ge=0.0, le=25.0)

tests/data/jsonschema/all_of_any_of_base_class_ref.json

@@ -0,0 +1,57 @@
+{
+  "$id": "https://example.com/schemas/MapState.json",
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "MapState",
+  "allOf": [
+    {
+      "anyOf": [
+        {
+          "type": "object",
+          "properties": {
+            "latitude": {"type": "number", "minimum": -90, "maximum": 90},
+            "longitude": {"type": "number", "minimum": -180, "maximum": 180},
+            "zoom": {"type": "number", "minimum": 0, "maximum": 25, "default": 0},
+            "bearing": {"type": "number"},
+            "pitch": {"type": "number", "minimum": 0, "exclusiveMaximum": 90},
+            "dragRotate": {"type": "boolean"},
+            "mapSplitMode": {"type": "string", "const": "SINGLE_MAP"},
+            "isSplit": {"type": "boolean", "const": false, "default": false}
+          },
+          "required": ["latitude", "longitude", "pitch", "mapSplitMode"]
+        },
+        {
+          "type": "object",
+          "properties": {
+            "latitude": {"$ref": "#/allOf/0/anyOf/0/properties/latitude"},
+            "longitude": {"$ref": "#/allOf/0/anyOf/0/properties/longitude"},
+            "zoom": {"$ref": "#/allOf/0/anyOf/0/properties/zoom"},
+            "bearing": {"$ref": "#/allOf/0/anyOf/0/properties/bearing"},
+            "pitch": {"$ref": "#/allOf/0/anyOf/0/properties/pitch"},
+            "dragRotate": {"$ref": "#/allOf/0/anyOf/0/properties/dragRotate"},
+            "mapSplitMode": {"type": "string", "const": "SWIPE_COMPARE"},
+            "isSplit": {"type": "boolean", "const": true, "default": true}
+          },
+          "required": ["latitude", "longitude", "pitch", "mapSplitMode"]
+        }
+      ]
+    },
+    {
+      "anyOf": [
+        {
+          "type": "object",
+          "properties": {
+            "mapViewMode": {"type": "string", "const": "MODE_2D"}
+          },
+          "required": ["mapViewMode"]
+        },
+        {
+          "type": "object",
+          "properties": {
+            "mapViewMode": {"type": "string", "const": "MODE_3D"}
+          },
+          "required": ["mapViewMode"]
+        }
+      ]
+    }
+  ]
+}

tests/main/jsonschema/test_main_jsonschema.py

@@ -2024,6 +2024,17 @@ def test_main_all_of_any_of(output_dir: Path) -> None:
     )
 
 
+def test_main_all_of_any_of_base_class_ref(output_file: Path) -> None:
+    """Test allOf/anyOf with base class references to avoid invalid imports."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "all_of_any_of_base_class_ref.json",
+        output_path=output_file,
+        input_file_type="jsonschema",
+        assert_func=assert_file_content,
+        extra_args=["--snake-case-field", "--use-double-quotes", "--reuse-model"],
+    )
+
+
 def test_main_all_of_one_of(output_dir: Path) -> None:
     """Test combination of allOf and oneOf."""
     run_main_and_assert(