Remove legacy annotation reassignment

Author: koxudaxi

src/datamodel_code_generator/model/pydantic_v2/base_model.py

@@ -105,6 +105,13 @@ def validate_min_max_items(cls, values: Any) -> dict[str, Any]:  # noqa: N805
 _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,
@@ -119,6 +126,22 @@ def _uses_legacy_pydantic_extra_template(template_file_path: str) -> bool:
     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:
@@ -555,6 +578,9 @@ def _render(self, *args: Any, **kwargs: Any) -> str:
                         "}",
                         *(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

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

@@ -17,7 +17,3 @@ class KnownAndExtra(BaseModel):
         '__pydantic_extra__': Dict[str, int],
     }
     name: str
-
-
-KnownAndExtra.__annotations__['__pydantic_extra__'] = Dict[str, int]
-KnownAndExtra.model_rebuild(force=True)

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

@@ -19,8 +19,4 @@ class Node(BaseModel):
     name: str
 
 
-Node.__annotations__['__pydantic_extra__'] = Dict[str, "Node"]
-Node.model_rebuild(force=True)
-
-
 Node.model_rebuild()

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
@@ -261,9 +264,34 @@ def test_pydantic_v2_legacy_extra_template_supports_relative_custom_path() -> No
     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(