Fix legacy pydantic extra templates

src/datamodel_code_generator/model/pydantic_v2/base_model.py

@@ -8,6 +8,7 @@
 
 import re
 from collections import defaultdict
+from functools import lru_cache
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Optional, cast
 
@@ -101,13 +102,46 @@ def validate_min_max_items(cls, values: Any) -> dict[str, Any]:  # noqa: N805
 _NO_ALIAS_INTERNAL_KEY = "_no_alias"
 _MISSING_SENTINEL = "MISSING"
 _CONFIG_ITEMS_TEMPLATE_DATA_KEY = "config_items"
+_LEGACY_PYDANTIC_EXTRA_TEMPLATE_PATTERN = re.compile(
+    r"{%-?\s*(?:if|elif)\s+(?:not\s+)?field\.use_pydantic_extra_annotation_assignment\b"
+)
+_LEGACY_PYDANTIC_EXTRA_POST_CLASS_PATTERN = re.compile(
+    r"(?m)^(?P<class_name>[^.\s]+)[ \t]*\.[ \t]*__annotations__[ \t]*"
+    r"\[[ \t]*(?P<quote>['\"])__pydantic_extra__(?P=quote)[ \t]*\][ \t]*=[^\r\n]*\r?\n"
+    r"(?:[ \t]*\r?\n)*"
+    r"(?P=class_name)[ \t]*\.[ \t]*model_rebuild[ \t]*"
+    r"\([ \t]*force[ \t]*=[ \t]*True[ \t]*\)[ \t]*(?:#[^\r\n]*)?(?:\r?\n)?"
+)
 _ALIAS_GENERATOR_IMPORTS: dict[str, Import] = {
     AliasGenerator.ToCamel.value: IMPORT_ALIAS_GENERATOR_TO_CAMEL,
     AliasGenerator.ToPascal.value: IMPORT_ALIAS_GENERATOR_TO_PASCAL,
     AliasGenerator.ToSnake.value: IMPORT_ALIAS_GENERATOR_TO_SNAKE,
 }
 
 
+@lru_cache(maxsize=16)
+def _uses_legacy_pydantic_extra_template(template_file_path: str) -> bool:
+    """Return whether a custom template uses the pre-0.68.1 typed-extra property."""
+    template_source = Path(template_file_path).read_text(encoding="utf-8")
+    return bool(_LEGACY_PYDANTIC_EXTRA_TEMPLATE_PATTERN.search(template_source))
+
+
+def _strip_legacy_pydantic_extra_post_class_assignment(rendered: str, class_name: str) -> str:
+    """Remove the unsupported post-class typed-extra assignment for the rendered model."""
+    if (
+        assignment := next(
+            (
+                match
+                for match in _LEGACY_PYDANTIC_EXTRA_POST_CLASS_PATTERN.finditer(rendered)
+                if match["class_name"] == class_name
+            ),
+            None,
+        )
+    ) is None:
+        return rendered
+    return f"{rendered[: assignment.start()]}{rendered[assignment.end() :]}"
+
+
 def _alias_generator_name(value: Any) -> str | None:
     generator_name: str | None = None
     match value:
@@ -324,6 +358,11 @@ def use_pydantic_extra_annotations_dict(self) -> bool:
         """Return whether typed extras need a class-body __annotations__ dict."""
         return self.is_pydantic_extra_field and not self.use_pydantic_extra_plain_annotation
 
+    @property
+    def use_pydantic_extra_annotation_assignment(self) -> bool:
+        """Support the typed-extra property used by pre-0.68.1 custom templates."""
+        return self.use_pydantic_extra_annotations_dict
+
     @property
     def pydantic_extra_type_hint(self) -> str:
         """Return a Dict-based type hint for Pydantic 2.0 typed extras."""
@@ -518,6 +557,32 @@ class BaseModel(BaseModelBase):
         ConfigAttribute("use_attribute_docstrings", "use_attribute_docstrings", False),  # noqa: FBT003
     ]
 
+    def _render(self, *args: Any, **kwargs: Any) -> str:
+        """Render evaluated typed extras for legacy custom templates."""
+        if self._custom_template_dir is None:
+            return super()._render(*args, **kwargs)
+
+        match self.template.filename:
+            case str() as filename if _uses_legacy_pydantic_extra_template(filename):
+                if field := next(
+                    (
+                        field
+                        for field in self.fields
+                        if isinstance(field, DataModelField) and field.use_pydantic_extra_annotations_dict
+                    ),
+                    None,
+                ):
+                    kwargs["class_body_lines"] = [
+                        "__annotations__ = {",
+                        f"    '__pydantic_extra__': {field.pydantic_extra_type_hint},",
+                        "}",
+                        *(kwargs.get("class_body_lines") or ()),
+                    ]
+                    return _strip_legacy_pydantic_extra_post_class_assignment(
+                        super()._render(*args, **kwargs), kwargs["class_name"]
+                    )
+        return super()._render(*args, **kwargs)
+
     @classmethod
     def render_module_code(cls, models: list[DataModel]) -> str:
         """Render shared schema runtime validation helpers for the module."""

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

@@ -0,0 +1,19 @@
+# generated by datamodel-codegen:
+#   filename:  additional_properties_schema_with_properties.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Dict
+
+from pydantic import BaseModel, ConfigDict
+
+
+class KnownAndExtra(BaseModel):
+    model_config = ConfigDict(
+        extra='allow',
+    )
+    __annotations__ = {
+        '__pydantic_extra__': Dict[str, int],
+    }
+    name: str

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

@@ -0,0 +1,22 @@
+# generated by datamodel-codegen:
+#   filename:  additional_properties_self_ref.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Dict
+
+from pydantic import BaseModel, ConfigDict
+
+
+class Node(BaseModel):
+    model_config = ConfigDict(
+        extra='allow',
+    )
+    __annotations__ = {
+        '__pydantic_extra__': Dict[str, "Node"],
+    }
+    name: str
+
+
+Node.model_rebuild()

tests/data/templates_pydantic_extra_pre_3593/pydantic_v2/BaseModel.jinja2

@@ -0,0 +1,71 @@
+{% for decorator in decorators -%}
+{{ decorator }}
+{% endfor -%}
+class {{ class_name }}({%- if schema_runtime_validation_use_base -%}{{ schema_runtime_validation_base_class_name }}{% if base_class != 'BaseModel' %}, {{ base_class }}{% endif %}{%- else -%}{{ base_class }}{%- endif -%}):{% if comment is defined %}  # {{ comment }}{% endif %}
+{%- if description %}
+    {{ description }}
+{%- endif %}
+{%- if not fields and not description and not config and not class_body_lines and not schema_runtime_validation %}
+    pass
+{%- endif %}
+{%- if config %}
+{%- filter indent(4) %}
+{% include 'ConfigDict.jinja2' %}
+{%- endfilter %}
+{%- endif %}
+{%- for line in class_body_lines %}
+    {{ line }}
+{%- endfor %}
+{%- include 'schema_runtime_validation.jinja2' %}
+{%- for field in fields %}
+    {%- if not field.use_pydantic_extra_annotation_assignment and not field.annotated and field.field %}
+    {{ field.name }}: {{ field.type_hint }} = {{ field.field }}
+    {%- elif not field.use_pydantic_extra_annotation_assignment %}
+    {%- if field.annotated %}
+    {{ field.name }}: {{ field.annotated }}
+    {%- else %}
+    {{ field.name }}: {{ field.type_hint }}
+    {%- endif %}
+    {%- if not field.has_default_factory_in_field and (not field.required or field.use_default_with_required) and (field.represented_default != 'None' or not field.strip_default_none or field.data_type.is_optional)
+            %} = {{ field.represented_default }}
+    {%- endif -%}
+    {%- endif %}
+    {%- if field.docstring %}
+    {{ field.docstring }}
+{%- if field.use_inline_field_description and not loop.last %}
+
+{% endif %}
+    {%- elif field.inline_field_docstring %}
+    {{ field.inline_field_docstring }}
+{%- if not loop.last %}
+
+{% endif %}
+    {%- endif %}
+{%- for method in methods -%}
+    {{ method }}
+{%- endfor -%}
+{%- endfor -%}
+{%- if prepared_validators %}
+{% for v in prepared_validators %}
+
+    @field_validator({{ v.fields_str }}, {{ v.mode_str }})
+    @classmethod
+{%- if v.mode == 'wrap' %}
+    def {{ v.method_name }}(cls, v: Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo) -> Any:
+        return {{ v.function_name }}(v, handler, info)
+{%- elif v.mode == 'plain' %}
+    def {{ v.method_name }}(cls, v: Any) -> Any:
+        return {{ v.function_name }}(v)
+{%- else %}
+    def {{ v.method_name }}(cls, v: Any, info: ValidationInfo) -> Any:
+        return {{ v.function_name }}(v, info)
+{%- endif %}
+{%- endfor %}
+{%- endif %}
+{%- for field in fields %}
+{%- if field.use_pydantic_extra_annotation_assignment %}
+
+{{ class_name }}.__annotations__['__pydantic_extra__'] = {{ field.pydantic_extra_type_hint }}
+{{ class_name }}.model_rebuild(force=True)
+{%- endif %}
+{%- endfor %}

tests/main/jsonschema/test_main_jsonschema.py

@@ -6056,6 +6056,119 @@ def test_main_jsonschema_additional_properties_schema_with_properties(output_fil
     )
 
 
+@pytest.mark.parametrize(
+    ("input_name", "expected_file", "disable_future_imports", "old_style_template"),
+    [
+        pytest.param(
+            "additional_properties_schema_with_properties.json",
+            "additional_properties_schema_with_properties_legacy_custom_template.py",
+            False,
+            False,
+            id="scalar-future",
+        ),
+        pytest.param(
+            "additional_properties_schema_with_properties.json",
+            "additional_properties_schema_with_properties_py313_no_future_imports.py",
+            True,
+            False,
+            id="scalar-no-future",
+        ),
+        pytest.param(
+            "additional_properties_self_ref.json",
+            "additional_properties_self_ref_legacy_custom_template.py",
+            False,
+            False,
+            id="self-ref-future",
+        ),
+        pytest.param(
+            "additional_properties_self_ref.json",
+            "additional_properties_self_ref_py314.py",
+            True,
+            False,
+            id="self-ref-no-future",
+        ),
+        pytest.param(
+            "additional_properties_schema_with_properties.json",
+            "additional_properties_schema_with_properties_legacy_custom_template.py",
+            False,
+            True,
+            id="scalar-future-old-style",
+        ),
+    ],
+)
+def test_main_jsonschema_legacy_pydantic_extra_custom_template(
+    input_name: str,
+    expected_file: str,
+    disable_future_imports: bool,
+    old_style_template: bool,
+    output_file: Path,
+    tmp_path: Path,
+) -> None:
+    """Test pre-0.68.1 custom templates keep typed-extra runtime validation."""
+    mode_args = ["--disable-future-imports"] if disable_future_imports else []
+    template_dir = (DATA_PATH / "templates_pydantic_extra_pre_3593").relative_to(Path.cwd())
+    copy_files = None
+    if old_style_template:
+        copy_files = [(template_dir / "pydantic_v2/BaseModel.jinja2", tmp_path / "BaseModel.jinja2")]
+        template_dir = tmp_path
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / input_name,
+        output_path=output_file,
+        input_file_type="jsonschema",
+        assert_func=assert_file_content,
+        expected_file=expected_file,
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+            "--custom-template-dir",
+            str(template_dir),
+            *mode_args,
+        ],
+        copy_files=copy_files,
+        force_exec_validation=True,
+    )
+
+    validation_case: tuple[str, str, str, str, tuple[str, ...], object] | None = None
+    match input_name:
+        case "additional_properties_schema_with_properties.json":
+            validation_case = (
+                "KnownAndExtra",
+                '{"name":"known","size":1}',
+                '{"name":"known","size":[]}',
+                "int_type",
+                ("__pydantic_extra__",),
+                {"size": 1},
+            )
+        case "additional_properties_self_ref.json":
+            validation_case = (
+                "Node",
+                '{"name":"root","child":{"name":"leaf"}}',
+                '{"name":"root","child":{"name":"leaf","bad":1}}',
+                "model_type",
+                ("__pydantic_extra__", "child", "name"),
+                "leaf",
+            )
+    if validation_case is None:  # pragma: no cover
+        raise AssertionError(input_name)
+    model_name, valid_json, invalid_json, expected_error_type, expected_attribute_path, expected_attribute_value = (
+        validation_case
+    )
+
+    assert_generated_model_json_validation(
+        output_file,
+        module_name=(
+            f"legacy_pydantic_extra_{Path(input_name).stem}_{'no_future' if disable_future_imports else 'future'}"
+            f"{'_old_style' if old_style_template else ''}"
+        ),
+        model_name=model_name,
+        valid_json=valid_json,
+        invalid_json=invalid_json,
+        expected_error_type=expected_error_type,
+        expected_attribute_path=expected_attribute_path,
+        expected_attribute_value=expected_attribute_value,
+    )
+
+
 @BLACK_PY314_SKIP
 def test_main_jsonschema_additional_properties_schema_with_properties_target_python_314(output_file: Path) -> None:
     """Test Python 3.14 target keeps typed extras as native deferred annotations."""

tests/model/test_base.py

@@ -41,6 +41,9 @@
 from datamodel_code_generator.model.pydantic_base import DataModelField as PydanticBaseDataModelField
 from datamodel_code_generator.model.pydantic_v2 import BaseModel
 from datamodel_code_generator.model.pydantic_v2 import DataModelField as PydanticV2DataModelField
+from datamodel_code_generator.model.pydantic_v2.base_model import (
+    _strip_legacy_pydantic_extra_post_class_assignment,
+)
 from datamodel_code_generator.model.pydantic_v2.imports import IMPORT_FIELD, IMPORT_MISSING
 from datamodel_code_generator.reference import Reference
 from datamodel_code_generator.types import ANY, NONE, DataType, Types
@@ -245,6 +248,50 @@ def test_pydantic_v2_extra_annotation_mode_uses_plain_annotation_for_native_defe
     assert IMPORT_DICT not in field.imports
 
 
+def test_pydantic_v2_legacy_extra_template_supports_relative_custom_path() -> None:
+    """Test legacy typed extras render through a relative custom template path."""
+    field = PydanticV2DataModelField(
+        name="__pydantic_extra__",
+        data_type=DataType(type="str", is_dict=True),
+        required=True,
+    )
+    model = BaseModel(
+        fields=[field],
+        reference=Reference(path="Model", original_name="Model", name="Model"),
+        custom_template_dir=Path("tests/data/templates_pydantic_extra_pre_3593"),
+    )
+
+    rendered = model.render()
+
+    assert "'__pydantic_extra__': Dict[str, str]," in rendered
+    assert "Model.__annotations__['__pydantic_extra__']" not in rendered
+    assert "Model.model_rebuild(force=True)" not in rendered
+    assert "locals()" not in rendered
+
+
+def test_strip_legacy_pydantic_extra_post_class_assignment_is_model_scoped() -> None:
+    """Strip only the target model's old assignment while preserving rebuilds and helpers."""
+    rendered = (
+        "Helper.__annotations__['__pydantic_extra__'] = Dict[str, int]\n"
+        "Helper.model_rebuild(force=True)\n"
+        'Model . __annotations__ [ "__pydantic_extra__" ] = Dict[str, int]\r\n'
+        "\r\n"
+        "Model . model_rebuild ( force = True )  # legacy\r\n"
+        "Model.model_rebuild()\n"
+    )
+
+    assert _strip_legacy_pydantic_extra_post_class_assignment(rendered, "Missing") == rendered
+    assert _strip_legacy_pydantic_extra_post_class_assignment(rendered, "Model") == (
+        "Helper.__annotations__['__pydantic_extra__'] = Dict[str, int]\n"
+        "Helper.model_rebuild(force=True)\n"
+        "Model.model_rebuild()\n"
+    )
+    assert not _strip_legacy_pydantic_extra_post_class_assignment(
+        "℘Model.__annotations__['__pydantic_extra__'] = Dict[str, int]\n℘Model.model_rebuild(force=True)\n",
+        "℘Model",
+    )
+
+
 def test_pydantic_v2_missing_sentinel_default_keeps_explicit_default() -> None:
     """Test explicit defaults are not replaced by the MISSING sentinel."""
     field = PydanticV2DataModelField(