Fix legacy pydantic extra templates

Author: koxudaxi

src/datamodel_code_generator/model/_custom_template.py

@@ -0,0 +1,124 @@
+"""Compatibility handling for custom Jinja templates."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from jinja2 import FileSystemLoader, nodes
+
+from datamodel_code_generator import Error
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any
+
+    from jinja2 import Environment
+
+_LEGACY_PYDANTIC_EXTRA_ATTRIBUTE_NAME = "use_pydantic_extra_annotation_assignment"
+_LEGACY_PYDANTIC_EXTRA_ATTRIBUTE = f"field.{_LEGACY_PYDANTIC_EXTRA_ATTRIBUTE_NAME}"
+_LEGACY_PYDANTIC_EXTRA_ATTRIBUTE_COUNT = 3
+_PYDANTIC_EXTRA_ANNOTATIONS_ATTRIBUTE_NAME = "use_pydantic_extra_annotations_dict"
+_PYDANTIC_EXTRA_ANNOTATIONS_ATTRIBUTE = f"field.{_PYDANTIC_EXTRA_ANNOTATIONS_ATTRIBUTE_NAME}"
+_PYDANTIC_V2_BASE_MODEL_TEMPLATE = "BaseModel.jinja2"
+_CLASS_BODY_LINES_BLOCK = """{%- for line in class_body_lines %}
+    {{ line }}
+{%- endfor %}"""
+_PYDANTIC_EXTRA_ANNOTATIONS_BLOCK = """{%- for field in fields %}
+{%- if field.use_pydantic_extra_annotations_dict %}
+    if (__annotations__ := locals().get('__annotations__')) is None:
+        __annotations__ = {}
+    __annotations__['__pydantic_extra__'] = {{ field.pydantic_extra_type_hint }}
+{%- endif %}
+{%- endfor %}
+"""
+_LEGACY_PYDANTIC_EXTRA_FIELD_IF = (
+    "{%- if not field.use_pydantic_extra_annotation_assignment and not field.annotated and field.field %}"
+)
+_PYDANTIC_EXTRA_FIELD_IF = (
+    "{%- if not field.use_pydantic_extra_annotations_dict and not field.annotated and field.field %}"
+)
+_LEGACY_PYDANTIC_EXTRA_FIELD_ELIF = "{%- elif not field.use_pydantic_extra_annotation_assignment %}"
+_PYDANTIC_EXTRA_FIELD_ELIF = "{%- elif not field.use_pydantic_extra_annotations_dict %}"
+_LEGACY_PYDANTIC_EXTRA_ASSIGNMENT_BLOCK = """
+{%- 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 %}"""
+
+
+def _legacy_template_error(filename: str) -> Error:
+    return Error(
+        "The custom Pydantic v2 BaseModel template uses the removed "
+        "'use_pydantic_extra_annotation_assignment' property, but its typed-extra blocks do not match "
+        f"the supported legacy template: {filename}. Copy the current pydantic_v2/BaseModel.jinja2 "
+        "template and reapply the custom changes."
+    )
+
+
+def _upgrade_legacy_pydantic_extra_template(environment: Environment, source: str, filename: str) -> str:
+    """Upgrade the pre-0.68.1 typed-extra blocks without changing other custom content."""
+    if not (legacy_attribute_count := source.count(_LEGACY_PYDANTIC_EXTRA_ATTRIBUTE)):
+        return source
+
+    replacements = (
+        (
+            _CLASS_BODY_LINES_BLOCK,
+            f"{_PYDANTIC_EXTRA_ANNOTATIONS_BLOCK}{_CLASS_BODY_LINES_BLOCK}",
+        ),
+        (_LEGACY_PYDANTIC_EXTRA_FIELD_IF, _PYDANTIC_EXTRA_FIELD_IF),
+        (_LEGACY_PYDANTIC_EXTRA_FIELD_ELIF, _PYDANTIC_EXTRA_FIELD_ELIF),
+        (_LEGACY_PYDANTIC_EXTRA_ASSIGNMENT_BLOCK, ""),
+    )
+    match tuple(source.count(old) for old, _ in replacements):
+        case (1, 1, 1, 1):
+            positions = tuple(source.index(old) for old, _ in replacements)
+        case _:
+            positions = ()
+
+    positions_are_ordered = bool(positions and positions == tuple(sorted(positions)))
+    if (
+        legacy_attribute_count != _LEGACY_PYDANTIC_EXTRA_ATTRIBUTE_COUNT
+        or source.count(_PYDANTIC_EXTRA_ANNOTATIONS_ATTRIBUTE)
+        or not positions_are_ordered
+    ):
+        parsed = environment.parse(source)
+        active_legacy_attributes = active_new_attributes = 0
+        for node in parsed.find_all(nodes.Getattr):
+            if not isinstance(node.node, nodes.Name) or node.node.name != "field":
+                continue
+            match node.attr:
+                case attribute if attribute == _LEGACY_PYDANTIC_EXTRA_ATTRIBUTE_NAME:
+                    active_legacy_attributes += 1
+                case attribute if attribute == _PYDANTIC_EXTRA_ANNOTATIONS_ATTRIBUTE_NAME:
+                    active_new_attributes += 1
+        if not active_legacy_attributes:
+            return source
+        if (
+            active_legacy_attributes != _LEGACY_PYDANTIC_EXTRA_ATTRIBUTE_COUNT
+            or active_new_attributes
+            or not positions_are_ordered
+        ):
+            raise _legacy_template_error(filename)
+
+    for old, new in replacements:
+        source = source.replace(old, new, 1)
+    return source
+
+
+class CustomTemplateFileSystemLoader(FileSystemLoader):
+    """Load custom templates with narrowly scoped legacy compatibility."""
+
+    def get_source(
+        self,
+        environment: Environment,
+        template: str,
+    ) -> tuple[str, str, Callable[..., Any]]:
+        source, filename, is_up_to_date = super().get_source(environment, template)
+        if template != _PYDANTIC_V2_BASE_MODEL_TEMPLATE:
+            return source, filename, is_up_to_date
+        if (upgraded_source := _upgrade_legacy_pydantic_extra_template(environment, source, filename)) is source:
+            return source, filename, is_up_to_date
+        return upgraded_source, filename, is_up_to_date

src/datamodel_code_generator/model/base.py

@@ -56,6 +56,7 @@
 })
 _MODULE_NAME_INVALID_CHAR_PATTERN = re.compile(r"[^0-9a-zA-Z_]")
 _MODULE_NAME_INVALID_CHAR_WITH_DOTS_PATTERN = re.compile(r"[^0-9a-zA-Z_.]")
+_PYDANTIC_V2_TEMPLATE_SUBDIR = "pydantic_v2"
 
 
 @dataclass(frozen=True, slots=True)
@@ -991,7 +992,14 @@ def _get_environment(template_subdir: Path, custom_template_dir: Path | None) ->
     if custom_template_dir is not None:
         custom_dir = custom_template_dir / template_subdir
         if cached_path_exists(custom_dir):
-            loaders.append(FileSystemLoader(str(custom_dir)))
+            if template_subdir.as_posix() == _PYDANTIC_V2_TEMPLATE_SUBDIR:
+                from datamodel_code_generator.model._custom_template import (  # noqa: PLC0415
+                    CustomTemplateFileSystemLoader,
+                )
+
+                loaders.append(CustomTemplateFileSystemLoader(str(custom_dir)))
+            else:
+                loaders.append(FileSystemLoader(str(custom_dir)))
             has_custom_loader = True
 
     loaders.append(FileSystemLoader(str(TEMPLATE_DIR / template_subdir)))
@@ -1021,8 +1029,18 @@ def _get_environment_with_absolute_path(absolute_template_dir: Path, builtin_sub
     """Get or create a cached Jinja2 Environment for absolute path templates."""
     from jinja2 import ChoiceLoader, FileSystemLoader  # noqa: PLC0415
 
+    custom_loader: FileSystemLoader
+    if builtin_subdir.as_posix() == _PYDANTIC_V2_TEMPLATE_SUBDIR:
+        from datamodel_code_generator.model._custom_template import (  # noqa: PLC0415
+            CustomTemplateFileSystemLoader,
+        )
+
+        custom_loader = CustomTemplateFileSystemLoader(str(absolute_template_dir))
+    else:
+        custom_loader = FileSystemLoader(str(absolute_template_dir))
+
     loaders: list[FileSystemLoader] = [
-        FileSystemLoader(str(absolute_template_dir)),
+        custom_loader,
         FileSystemLoader(str(TEMPLATE_DIR / builtin_subdir)),
     ]
     return _build_environment(ChoiceLoader(loaders))

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',
+    )
+    if (__annotations__ := locals().get('__annotations__')) is None:
+        __annotations__ = {}
+    __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',
+    )
+    if (__annotations__ := locals().get('__annotations__')) is None:
+        __annotations__ = {}
+    __annotations__['__pydantic_extra__'] = Dict[str, "Node"]
+    name: str
+
+
+Node.model_rebuild()

tests/data/templates_extensions/pydantic_v2/BaseModel.jinja2

@@ -6,6 +6,7 @@ an alias for Bar: every pydantic model class consumes considerable memory. #}
 
 {% else -%}
 
+{# field.use_pydantic_extra_annotation_assignment is intentionally only a comment. #}
 {% for decorator in decorators -%}
 {{ decorator }}
 {% endfor -%}

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/data/templates_pydantic_extra_unsupported/pydantic_v2/BaseModel.jinja2

@@ -0,0 +1 @@
+{{ field.use_pydantic_extra_annotation_assignment }}