Skip to content

Commit bed2766

Browse files
committed
Refactor pydantic v2 config generation
1 parent 2d379eb commit bed2766

9 files changed

Lines changed: 267 additions & 163 deletions

File tree

src/datamodel_code_generator/model/msgspec.py

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,6 @@
99
from math import isfinite
1010
from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypeVar
1111

12-
from pydantic import Field
13-
1412
from datamodel_code_generator.imports import IMPORT_UNION, Import
1513
from datamodel_code_generator.model import DataModel, DataModelFieldBase, _rebuild_model_with_datamodel_namespace
1614
from datamodel_code_generator.model.base import UNDEFINED, BaseClassDataType, _nested_model_default_factory
@@ -23,7 +21,7 @@
2321
IMPORT_MSGSPEC_UNSETTYPE,
2422
)
2523
from datamodel_code_generator.model.pydantic_base import (
26-
Constraints as _Constraints,
24+
PatternConstraints as _Constraints,
2725
)
2826
from datamodel_code_generator.model.type_alias import TypeAliasBase
2927
from datamodel_code_generator.model.types import DataTypeManager as _DataTypeManager
@@ -206,10 +204,6 @@ def create_base_class_model(
206204
class Constraints(_Constraints):
207205
"""Constraint model for msgspec fields."""
208206

209-
# To override existing pattern alias
210-
regex: Optional[str] = Field(None, alias="regex") # noqa: UP045
211-
pattern: Optional[str] = Field(None, alias="pattern") # noqa: UP045
212-
213207

214208
@lru_cache
215209
def get_neither_required_nor_nullable_type(type_: str, use_union_operator: bool) -> str: # noqa: FBT001

src/datamodel_code_generator/model/pydantic_base.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,11 @@ class Constraints(ConstraintsBase):
5252
regex: Optional[str] = Field(None, alias="pattern") # noqa: UP045
5353

5454

55+
class PatternConstraints(Constraints): # noqa: D101
56+
regex: Optional[str] = Field(None, alias="regex") # noqa: UP045
57+
pattern: Optional[str] = Field(None, alias="pattern") # noqa: UP045
58+
59+
5560
class DataModelField(DataModelFieldBase):
5661
"""Field implementation for Pydantic models."""
5762

src/datamodel_code_generator/model/pydantic_v2/__init__.py

Lines changed: 2 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,11 @@
66

77
from __future__ import annotations
88

9-
from typing import TYPE_CHECKING, Any, Dict, Optional # noqa: UP035
10-
11-
from pydantic import BaseModel as _BaseModel
9+
from typing import TYPE_CHECKING
1210

1311
from datamodel_code_generator.enums import UnionMode
1412

13+
from ._config import ConfigDict as ConfigDict
1514
from .base_model import BaseModel, DataModelField
1615
from .root_model import RootModel
1716
from .root_model_type_alias import RootModelTypeAlias
@@ -26,31 +25,6 @@ def dump_resolve_reference_action(class_names: Iterable[str]) -> str:
2625
return "\n".join(f"{class_name}.model_rebuild()" for class_name in class_names)
2726

2827

29-
class ConfigDict(_BaseModel):
30-
"""Pydantic v2 model_config options."""
31-
32-
extra: Optional[str] = None # noqa: UP045
33-
title: Optional[str] = None # noqa: UP045
34-
populate_by_name: Optional[bool] = None # noqa: UP045 # deprecated in v2.11+
35-
validate_by_name: Optional[bool] = None # noqa: UP045 # v2.11+
36-
validate_by_alias: Optional[bool] = None # noqa: UP045 # v2.11+
37-
allow_extra_fields: Optional[bool] = None # noqa: UP045
38-
extra_fields: Optional[str] = None # noqa: UP045
39-
from_attributes: Optional[bool] = None # noqa: UP045
40-
frozen: Optional[bool] = None # noqa: UP045
41-
arbitrary_types_allowed: Optional[bool] = None # noqa: UP045
42-
protected_namespaces: Optional[tuple[str, ...]] = None # noqa: UP045
43-
regex_engine: Optional[str] = None # noqa: UP045
44-
use_enum_values: Optional[bool] = None # noqa: UP045
45-
coerce_numbers_to_str: Optional[bool] = None # noqa: UP045
46-
use_attribute_docstrings: Optional[bool] = None # noqa: UP045
47-
json_schema_extra: Optional[Dict[str, Any]] = None # noqa: UP006, UP045
48-
49-
def dict(self, **kwargs: Any) -> dict[str, Any]: # type: ignore[override]
50-
"""Return dict for templates."""
51-
return self.model_dump(**kwargs)
52-
53-
5428
__all__ = [
5529
"BaseModel",
5630
"DataModelField",
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
"""Shared Pydantic v2 ConfigDict generation helpers."""
2+
3+
from __future__ import annotations
4+
5+
from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Optional # noqa: UP035
6+
7+
from pydantic import BaseModel as _BaseModel
8+
9+
from datamodel_code_generator.enums import TargetPydanticVersion
10+
11+
if TYPE_CHECKING:
12+
from collections.abc import Iterable
13+
14+
from datamodel_code_generator.types import DataType
15+
16+
17+
class ConfigAttribute(NamedTuple):
18+
"""Configuration attribute mapping for ConfigDict conversion."""
19+
20+
from_: str
21+
to: str
22+
invert: bool
23+
24+
25+
def get_config_extra(extra_template_data: dict[str, Any]) -> str | None:
26+
"""Get extra field configuration for ConfigDict."""
27+
additional_properties = extra_template_data.get("additionalProperties")
28+
unevaluated_properties = extra_template_data.get("unevaluatedProperties")
29+
allow_extra_fields = extra_template_data.get("allow_extra_fields")
30+
extra_fields = extra_template_data.get("extra_fields")
31+
32+
config_extra = None
33+
if allow_extra_fields or extra_fields == "allow":
34+
config_extra = "'allow'"
35+
elif extra_fields == "forbid":
36+
config_extra = "'forbid'"
37+
elif extra_fields == "ignore":
38+
config_extra = "'ignore'"
39+
elif additional_properties is True:
40+
config_extra = "'allow'"
41+
elif additional_properties is False:
42+
config_extra = "'forbid'"
43+
elif unevaluated_properties is True:
44+
config_extra = "'allow'"
45+
elif unevaluated_properties is False:
46+
config_extra = "'forbid'"
47+
return config_extra
48+
49+
50+
def get_config_attributes(
51+
extra_template_data: dict[str, Any],
52+
*,
53+
config_attributes_v2: list[ConfigAttribute],
54+
config_attributes_v2_11: list[ConfigAttribute],
55+
) -> list[ConfigAttribute]:
56+
"""Get config attributes based on target Pydantic version."""
57+
target_version = extra_template_data.get("target_pydantic_version")
58+
if target_version == TargetPydanticVersion.V2_11:
59+
return config_attributes_v2_11
60+
return config_attributes_v2
61+
62+
63+
def build_base_config_parameters(
64+
*,
65+
extra_template_data: dict[str, Any],
66+
all_data_types: Iterable[DataType],
67+
config_attributes_v2: list[ConfigAttribute],
68+
config_attributes_v2_11: list[ConfigAttribute],
69+
include_extra: bool = True,
70+
) -> dict[str, Any]:
71+
"""Build shared ConfigDict parameters for Pydantic v2 models."""
72+
config_parameters: dict[str, Any] = {}
73+
74+
if include_extra and (extra := get_config_extra(extra_template_data)):
75+
config_parameters["extra"] = extra
76+
77+
config_attributes = get_config_attributes(
78+
extra_template_data,
79+
config_attributes_v2=config_attributes_v2,
80+
config_attributes_v2_11=config_attributes_v2_11,
81+
)
82+
for from_, to, invert in config_attributes:
83+
if from_ in extra_template_data:
84+
config_parameters[to] = not extra_template_data[from_] if invert else extra_template_data[from_]
85+
86+
for data_type in all_data_types:
87+
if data_type.is_custom_type: # pragma: no cover
88+
config_parameters["arbitrary_types_allowed"] = True
89+
break
90+
91+
return config_parameters
92+
93+
94+
class ConfigDict(_BaseModel):
95+
"""Pydantic v2 model_config options."""
96+
97+
extra: Optional[str] = None # noqa: UP045
98+
title: Optional[str] = None # noqa: UP045
99+
populate_by_name: Optional[bool] = None # noqa: UP045 # deprecated in v2.11+
100+
validate_by_name: Optional[bool] = None # noqa: UP045 # v2.11+
101+
validate_by_alias: Optional[bool] = None # noqa: UP045 # v2.11+
102+
allow_extra_fields: Optional[bool] = None # noqa: UP045
103+
extra_fields: Optional[str] = None # noqa: UP045
104+
from_attributes: Optional[bool] = None # noqa: UP045
105+
frozen: Optional[bool] = None # noqa: UP045
106+
arbitrary_types_allowed: Optional[bool] = None # noqa: UP045
107+
protected_namespaces: Optional[tuple[str, ...]] = None # noqa: UP045
108+
regex_engine: Optional[str] = None # noqa: UP045
109+
use_enum_values: Optional[bool] = None # noqa: UP045
110+
coerce_numbers_to_str: Optional[bool] = None # noqa: UP045
111+
use_attribute_docstrings: Optional[bool] = None # noqa: UP045
112+
json_schema_extra: Optional[Dict[str, Any]] = None # noqa: UP006, UP045
113+
114+
def dict(self, **kwargs: Any) -> dict[str, Any]: # type: ignore[override]
115+
"""Return dict for templates."""
116+
return self.model_dump(**kwargs)

src/datamodel_code_generator/model/pydantic_v2/base_model.py

Lines changed: 18 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,11 @@
88

99
import re
1010
from collections import defaultdict
11-
from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, Optional
11+
from typing import TYPE_CHECKING, Any, ClassVar, Optional
1212

13-
from pydantic import Field, ValidationError, field_validator, model_validator
13+
from pydantic import ValidationError, field_validator, model_validator
1414

1515
from datamodel_code_generator import Error
16-
from datamodel_code_generator.enums import TargetPydanticVersion
1716
from datamodel_code_generator.imports import IMPORT_ANY, IMPORT_DICT, Import
1817
from datamodel_code_generator.model import _rebuild_model_with_datamodel_namespace
1918
from datamodel_code_generator.model.base import (
@@ -23,14 +22,17 @@
2322
DataModelFieldBase,
2423
)
2524
from datamodel_code_generator.model.imports import IMPORT_CLASSVAR
25+
from datamodel_code_generator.model.pydantic_base import BaseModelBase
2626
from datamodel_code_generator.model.pydantic_base import (
27-
BaseModelBase,
27+
DataModelField as _PydanticBaseDataModelField,
2828
)
2929
from datamodel_code_generator.model.pydantic_base import (
30-
Constraints as _Constraints,
30+
PatternConstraints as _Constraints,
3131
)
32-
from datamodel_code_generator.model.pydantic_base import (
33-
DataModelField as _PydanticBaseDataModelField,
32+
from datamodel_code_generator.model.pydantic_v2._config import (
33+
ConfigAttribute,
34+
ConfigDict,
35+
build_base_config_parameters,
3436
)
3537
from datamodel_code_generator.model.pydantic_v2.imports import (
3638
IMPORT_ALIAS_CHOICES,
@@ -68,10 +70,6 @@ def __repr__(self) -> str:
6870
class Constraints(_Constraints):
6971
"""Pydantic v2 field constraints with pattern support."""
7072

71-
# To override existing pattern alias
72-
regex: Optional[str] = Field(None, alias="regex") # noqa: UP045
73-
pattern: Optional[str] = Field(None, alias="pattern") # noqa: UP045
74-
7573
@model_validator(mode="before") # ty: ignore
7674
def validate_min_max_items(cls, values: Any) -> dict[str, Any]: # noqa: N805
7775
"""Validate and convert minItems/maxItems to minLength/maxLength."""
@@ -273,14 +271,6 @@ def imports(self) -> tuple[Import, ...]:
273271
return base_imports
274272

275273

276-
class ConfigAttribute(NamedTuple):
277-
"""Configuration attribute mapping for ConfigDict conversion."""
278-
279-
from_: str
280-
to: str
281-
invert: bool
282-
283-
284274
class BaseModel(BaseModelBase):
285275
"""Pydantic v2 BaseModel with ConfigDict and pattern-based regex_engine support."""
286276

@@ -290,6 +280,7 @@ class BaseModel(BaseModelBase):
290280
BASE_CLASS_ALIAS: ClassVar[str] = "_BaseModel"
291281
SUPPORTS_DISCRIMINATOR: ClassVar[bool] = True
292282
SUPPORTS_FIELD_RENAMING: ClassVar[bool] = True
283+
SUPPORTS_CONFIG_EXTRA: ClassVar[bool] = True
293284
TYPED_EXTRA_FIELD_NAME: ClassVar[str] = "__pydantic_extra__"
294285
# In Pydantic 2.11+, populate_by_name is deprecated in favor of validate_by_name + validate_by_alias
295286
# Default to V2 compatible (populate_by_name) unless target_pydantic_version is specified
@@ -356,30 +347,20 @@ def __init__( # noqa: PLR0913
356347
keyword_only=keyword_only,
357348
treat_dot_as_module=treat_dot_as_module,
358349
)
359-
config_parameters: dict[str, Any] = {}
360-
361-
extra = self._get_config_extra()
362-
if extra:
363-
config_parameters["extra"] = extra
364-
365-
# Select CONFIG_ATTRIBUTES based on target_pydantic_version
366-
config_attributes = self._get_config_attributes()
367-
for from_, to, invert in config_attributes:
368-
if from_ in self.extra_template_data:
369-
config_parameters[to] = (
370-
not self.extra_template_data[from_] if invert else self.extra_template_data[from_]
371-
)
372-
for data_type in self.all_data_types:
373-
if data_type.is_custom_type: # pragma: no cover
374-
config_parameters["arbitrary_types_allowed"] = True
375-
break
350+
config_parameters = build_base_config_parameters(
351+
extra_template_data=self.extra_template_data,
352+
all_data_types=self.all_data_types,
353+
config_attributes_v2=self._CONFIG_ATTRIBUTES_V2,
354+
config_attributes_v2_11=self._CONFIG_ATTRIBUTES_V2_11,
355+
include_extra=self.SUPPORTS_CONFIG_EXTRA,
356+
)
376357

377358
if self._has_lookaround_pattern():
378359
config_parameters["regex_engine"] = '"python-re"'
379360

380361
if isinstance(self.extra_template_data.get("config"), dict):
381362
for key, value in self.extra_template_data["config"].items():
382-
config_parameters[key] = value # noqa: PERF403
363+
config_parameters[key] = value
383364

384365
# Handle json_schema_extra from schema extensions (x-* fields)
385366
model_extras = self.extra_template_data.get("model_extras")
@@ -388,47 +369,11 @@ def __init__( # noqa: PLR0913
388369
config_parameters["json_schema_extra"] = {**existing, **model_extras}
389370

390371
if config_parameters:
391-
from datamodel_code_generator.model.pydantic_v2 import ConfigDict # noqa: PLC0415
392-
393372
self.extra_template_data["config"] = ConfigDict.model_validate(config_parameters) # ty: ignore
394373
self._additional_imports.append(IMPORT_CONFIG_DICT)
395374

396375
self._process_validators()
397376

398-
def _get_config_extra(self) -> Literal["'allow'", "'forbid'", "'ignore'"] | None:
399-
additional_properties = self.extra_template_data.get("additionalProperties")
400-
unevaluated_properties = self.extra_template_data.get("unevaluatedProperties")
401-
allow_extra_fields = self.extra_template_data.get("allow_extra_fields")
402-
extra_fields = self.extra_template_data.get("extra_fields")
403-
404-
config_extra = None
405-
if allow_extra_fields or extra_fields == "allow":
406-
config_extra = "'allow'"
407-
elif extra_fields == "forbid":
408-
config_extra = "'forbid'"
409-
elif extra_fields == "ignore":
410-
config_extra = "'ignore'"
411-
elif additional_properties is True:
412-
config_extra = "'allow'"
413-
elif additional_properties is False:
414-
config_extra = "'forbid'"
415-
elif unevaluated_properties is True:
416-
config_extra = "'allow'"
417-
elif unevaluated_properties is False:
418-
config_extra = "'forbid'"
419-
return config_extra
420-
421-
def _get_config_attributes(self) -> list[ConfigAttribute]:
422-
"""Get config attributes based on target Pydantic version.
423-
424-
If target_pydantic_version is V2_11, use validate_by_name.
425-
Otherwise (V2 or not specified), use populate_by_name for compatibility.
426-
"""
427-
target_version = self.extra_template_data.get("target_pydantic_version")
428-
if target_version == TargetPydanticVersion.V2_11:
429-
return self._CONFIG_ATTRIBUTES_V2_11
430-
return self._CONFIG_ATTRIBUTES_V2
431-
432377
def _has_lookaround_pattern(self) -> bool:
433378
"""Check if any field has a regex pattern with lookaround assertions."""
434379
lookaround_regex = re.compile(r"\(\?<?[=!]")

0 commit comments

Comments
 (0)