Restore _try_rebuild_model and test_input_model_config_class after merge

Author: koxudaxi

src/datamodel_code_generator/__main__.py

@@ -1070,6 +1070,37 @@ def _filter_defs_by_strategy(
     return {**schema, "$defs": new_defs}
 
 
+def _try_rebuild_model(obj: type) -> None:
+    """Try to rebuild a Pydantic model, handling config models specially."""
+    module = getattr(obj, "__module__", "")
+    class_name = getattr(obj, "__name__", "")
+    config_classes = {"GenerateConfig", "ParserConfig", "ParseConfig"}
+    if module in {"datamodel_code_generator.config", "config"} and class_name in config_classes:
+        from datamodel_code_generator.model.base import DataModel, DataModelFieldBase  # noqa: PLC0415
+        from datamodel_code_generator.types import DataTypeManager, StrictTypes  # noqa: PLC0415
+
+        try:
+            from datamodel_code_generator.model.pydantic_v2 import UnionMode  # noqa: PLC0415
+        except ImportError:  # pragma: no cover
+            from typing import Any  # noqa: PLC0415
+
+            runtime_union_mode = Any
+        else:
+            runtime_union_mode = UnionMode
+
+        types_namespace = {
+            "Path": Path,
+            "DataModel": DataModel,
+            "DataModelFieldBase": DataModelFieldBase,
+            "DataTypeManager": DataTypeManager,
+            "StrictTypes": StrictTypes,
+            "UnionMode": runtime_union_mode,
+        }
+        obj.model_rebuild(_types_namespace=types_namespace)
+    else:
+        obj.model_rebuild()
+
+
 def _load_model_schema(  # noqa: PLR0912, PLR0914, PLR0915
     input_model: str,
     input_file_type: InputFileType,
@@ -1151,6 +1182,8 @@ def _load_model_schema(  # noqa: PLR0912, PLR0914, PLR0915
         if not hasattr(obj, "model_json_schema"):
             msg = "--input-model with Pydantic model requires Pydantic v2 runtime. Please upgrade Pydantic to v2."
             raise Error(msg)
+        if hasattr(obj, "model_rebuild"):  # pragma: no branch
+            _try_rebuild_model(obj)
         schema_generator = _get_input_model_json_schema_class()
         schema = obj.model_json_schema(schema_generator=schema_generator)
         schema = _add_python_type_for_unserializable(schema, obj)

tests/test_input_model.py

@@ -962,3 +962,14 @@ def test_input_model_ref_strategy_typeddict_reuse_foreign(tmp_path: Path) -> Non
             "class Profile",
         ],
     )
+
+
+@SKIP_PYDANTIC_V1
+def test_input_model_config_class(tmp_path: Path) -> None:
+    """Test that config classes like GenerateConfig are properly handled."""
+    run_input_model_and_assert(
+        input_model="datamodel_code_generator.config:GenerateConfig",
+        output_path=tmp_path / "output.py",
+        extra_args=["--output-model-type", "typing.TypedDict"],
+        expected_output_contains=["TypedDict", "Callable[[str], str]"],
+    )