Enhance TypeAlias handling to maintain model order with forward references (#2560)

Author: koxudaxi

src/datamodel_code_generator/parser/base.py

@@ -47,6 +47,7 @@
     DataModelFieldBase,
 )
 from datamodel_code_generator.model.enum import Enum, Member
+from datamodel_code_generator.model.type_alias import TypeAliasBase
 from datamodel_code_generator.parser import DefaultPutDict, LiteralType
 from datamodel_code_generator.reference import ModelResolver, Reference
 from datamodel_code_generator.types import DataType, DataTypeManager, StrictTypes
@@ -1174,22 +1175,28 @@ def __sort_models(
         models.sort(key=lambda x: x.class_name)
 
         imported = {i for v in imports.values() for i in v}
-        model_class_name_baseclasses: dict[DataModel, tuple[str, set[str]]] = {}
+        model_class_name_refs: dict[DataModel, tuple[str, set[str]]] = {}
         for model in models:
             class_name = model.class_name
-            model_class_name_baseclasses[model] = (
-                class_name,
-                {b.type_hint for b in model.base_classes if b.reference} - {class_name},
-            )
+            base_class_refs = {b.type_hint for b in model.base_classes if b.reference}
+            if base_class_refs:
+                refs = base_class_refs - {class_name}
+            elif isinstance(model, TypeAliasBase):
+                refs = {
+                    t.reference.short_name for f in model.fields for t in f.data_type.all_data_types if t.reference
+                } - {class_name}
+            else:
+                refs = set()
+            model_class_name_refs[model] = (class_name, refs)
 
         changed: bool = True
         while changed:
             changed = False
             resolved = imported.copy()
             for i in range(len(models) - 1):
                 model = models[i]
-                class_name, baseclasses = model_class_name_baseclasses[model]
-                if not baseclasses - resolved:
+                class_name, refs = model_class_name_refs[model]
+                if not refs - resolved:
                     resolved.add(class_name)
                     continue
                 models[i], models[i + 1] = models[i + 1], model

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

@@ -0,0 +1,25 @@
+# generated by datamodel-codegen:
+#   filename:  type_alias_forward_ref.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Literal, TypedDict
+
+from typing_extensions import TypeAliasType
+
+
+class BlobPart(TypedDict):
+    type: Literal['blob']
+    data: str
+
+
+FieldPlaceholder = TypeAliasType("FieldPlaceholder", None)
+
+
+class TextPart(TypedDict):
+    type: Literal['text']
+    content: str
+
+
+SystemInstructions = TypeAliasType("SystemInstructions", list[TextPart | BlobPart])

tests/data/jsonschema/type_alias_forward_ref.json

@@ -0,0 +1,45 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "_Placeholder",
+  "type": "null",
+  "$defs": {
+    "TextPart": {
+      "properties": {
+        "type": {
+          "const": "text",
+          "type": "string"
+        },
+        "content": {
+          "type": "string"
+        }
+      },
+      "required": ["type", "content"],
+      "title": "TextPart",
+      "type": "object"
+    },
+    "BlobPart": {
+      "properties": {
+        "type": {
+          "const": "blob",
+          "type": "string"
+        },
+        "data": {
+          "type": "string"
+        }
+      },
+      "required": ["type", "data"],
+      "title": "BlobPart",
+      "type": "object"
+    },
+    "SystemInstructions": {
+      "type": "array",
+      "items": {
+        "anyOf": [
+          {"$ref": "#/$defs/TextPart"},
+          {"$ref": "#/$defs/BlobPart"}
+        ]
+      },
+      "title": "SystemInstructions"
+    }
+  }
+}

tests/main/jsonschema/test_main_jsonschema.py

@@ -64,6 +64,27 @@ def test_main_inheritance_forward_ref_keep_model_order(output_file: Path, tmp_pa
     )
 
 
+@pytest.mark.benchmark
+def test_main_type_alias_forward_ref_keep_model_order(output_file: Path) -> None:
+    """Test TypeAliasType with forward references keeping model order."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "type_alias_forward_ref.json",
+        output_path=output_file,
+        input_file_type=None,
+        assert_func=assert_file_content,
+        extra_args=[
+            "--keep-model-order",
+            "--output-model-type",
+            "typing.TypedDict",
+            "--use-standard-collections",
+            "--use-union-operator",
+            "--use-type-alias",
+            "--target-python-version",
+            "3.10",
+        ],
+    )
+
+
 @pytest.mark.skip(reason="pytest-xdist does not support the test")
 def test_main_without_arguments() -> None:
     """Test main function without arguments raises SystemExit."""