Skip to content

Commit 276caac

Browse files
koxudaxikoxudaxi
authored
Deduplicate Pydantic v2 config helpers (#3350)
* Refactor pydantic v2 config generation * Preserve Pydantic v2 ConfigDict surface * Avoid copying Pydantic v2 config parameters * Cover RootModel ignored extra modes * Restore Pydantic v2 constraints surface * Avoid RootModel config scans
1 parent 577d495 commit 276caac

8 files changed

Lines changed: 277 additions & 133 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/_config.py‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
"""Shared Pydantic v2 ConfigDict generation helpers."""
2+
3+
from __future__ import annotations
4+
5+
from typing import TYPE_CHECKING, Any, Literal, NamedTuple
6+
7+
from datamodel_code_generator.enums import TargetPydanticVersion
8+
9+
if TYPE_CHECKING:
10+
from collections.abc import Iterable
11+
12+
from datamodel_code_generator.types import DataType
13+
14+
15+
class ConfigAttribute(NamedTuple):
16+
"""Configuration attribute mapping for ConfigDict conversion."""
17+
18+
from_: str
19+
to: str
20+
invert: bool
21+
22+
23+
ConfigExtra = Literal["'allow'", "'forbid'", "'ignore'"]
24+
ConfigParameterValue = bool | ConfigExtra | Literal['"python-re"'] | dict[str, Any]
25+
_CONFIG_EXTRA_KEYS: frozenset[str] = frozenset({
26+
"additionalProperties",
27+
"allow_extra_fields",
28+
"extra_fields",
29+
"unevaluatedProperties",
30+
})
31+
32+
33+
def get_config_extra(extra_template_data: dict[str, Any]) -> ConfigExtra | None:
34+
"""Get extra field configuration for ConfigDict."""
35+
additional_properties = extra_template_data.get("additionalProperties")
36+
unevaluated_properties = extra_template_data.get("unevaluatedProperties")
37+
allow_extra_fields = extra_template_data.get("allow_extra_fields")
38+
extra_fields = extra_template_data.get("extra_fields")
39+
40+
config_extra: ConfigExtra | None = None
41+
if allow_extra_fields or extra_fields == "allow":
42+
config_extra = "'allow'"
43+
elif extra_fields == "forbid":
44+
config_extra = "'forbid'"
45+
elif extra_fields == "ignore":
46+
config_extra = "'ignore'"
47+
elif additional_properties is True:
48+
config_extra = "'allow'"
49+
elif additional_properties is False:
50+
config_extra = "'forbid'"
51+
elif unevaluated_properties is True:
52+
config_extra = "'allow'"
53+
elif unevaluated_properties is False:
54+
config_extra = "'forbid'"
55+
return config_extra
56+
57+
58+
def get_config_attributes(
59+
extra_template_data: dict[str, Any],
60+
*,
61+
config_attributes_v2: list[ConfigAttribute],
62+
config_attributes_v2_11: list[ConfigAttribute],
63+
) -> list[ConfigAttribute]:
64+
"""Get config attributes based on target Pydantic version."""
65+
target_version = extra_template_data.get("target_pydantic_version")
66+
if target_version == TargetPydanticVersion.V2_11:
67+
return config_attributes_v2_11
68+
return config_attributes_v2
69+
70+
71+
def build_base_config_parameters(
72+
*,
73+
extra_template_data: dict[str, Any],
74+
all_data_types: Iterable[DataType],
75+
config_attributes_v2: list[ConfigAttribute],
76+
config_attributes_v2_11: list[ConfigAttribute],
77+
include_extra: bool = True,
78+
) -> dict[str, ConfigParameterValue]:
79+
"""Build shared ConfigDict parameters for Pydantic v2 models."""
80+
config_attributes = get_config_attributes(
81+
extra_template_data,
82+
config_attributes_v2=config_attributes_v2,
83+
config_attributes_v2_11=config_attributes_v2_11,
84+
)
85+
if (
86+
all_data_types == ()
87+
and not (include_extra and not _CONFIG_EXTRA_KEYS.isdisjoint(extra_template_data))
88+
and not any(from_ in extra_template_data for from_, _, _ in config_attributes)
89+
):
90+
return {}
91+
92+
config_parameters: dict[str, ConfigParameterValue] = {}
93+
94+
if include_extra and (extra := get_config_extra(extra_template_data)):
95+
config_parameters["extra"] = extra
96+
97+
for from_, to, invert in config_attributes:
98+
if from_ in extra_template_data:
99+
value: bool = extra_template_data[from_]
100+
config_parameters[to] = not value if invert else value
101+
102+
for data_type in all_data_types:
103+
if data_type.is_custom_type: # pragma: no cover
104+
config_parameters["arbitrary_types_allowed"] = True
105+
break
106+
107+
return config_parameters

‎src/datamodel_code_generator/model/pydantic_v2/base_model.py‎

Lines changed: 20 additions & 71 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, cast
1212

1313
from pydantic import Field, 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,15 +22,15 @@
2322
DataModelFieldBase,
2423
)
2524
from datamodel_code_generator.model.imports import IMPORT_CLASSVAR
26-
from datamodel_code_generator.model.pydantic_base import (
27-
BaseModelBase,
28-
)
29-
from datamodel_code_generator.model.pydantic_base import (
30-
Constraints as _Constraints,
31-
)
25+
from datamodel_code_generator.model.pydantic_base import BaseModelBase
26+
from datamodel_code_generator.model.pydantic_base import Constraints as _Constraints
3227
from datamodel_code_generator.model.pydantic_base import (
3328
DataModelField as _PydanticBaseDataModelField,
3429
)
30+
from datamodel_code_generator.model.pydantic_v2._config import (
31+
ConfigAttribute,
32+
build_base_config_parameters,
33+
)
3534
from datamodel_code_generator.model.pydantic_v2.imports import (
3635
IMPORT_ALIAS_CHOICES,
3736
IMPORT_BASE_MODEL,
@@ -273,14 +272,6 @@ def imports(self) -> tuple[Import, ...]:
273272
return base_imports
274273

275274

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

@@ -290,6 +281,8 @@ class BaseModel(BaseModelBase):
290281
BASE_CLASS_ALIAS: ClassVar[str] = "_BaseModel"
291282
SUPPORTS_DISCRIMINATOR: ClassVar[bool] = True
292283
SUPPORTS_FIELD_RENAMING: ClassVar[bool] = True
284+
SUPPORTS_CONFIG_EXTRA: ClassVar[bool] = True
285+
SUPPORTS_ARBITRARY_TYPES_ALLOWED: ClassVar[bool] = True
293286
TYPED_EXTRA_FIELD_NAME: ClassVar[str] = "__pydantic_extra__"
294287
# In Pydantic 2.11+, populate_by_name is deprecated in favor of validate_by_name + validate_by_alias
295288
# Default to V2 compatible (populate_by_name) unless target_pydantic_version is specified
@@ -356,35 +349,25 @@ def __init__( # noqa: PLR0913
356349
keyword_only=keyword_only,
357350
treat_dot_as_module=treat_dot_as_module,
358351
)
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
352+
config_parameters = build_base_config_parameters(
353+
extra_template_data=self.extra_template_data,
354+
all_data_types=self.all_data_types if self.SUPPORTS_ARBITRARY_TYPES_ALLOWED else (),
355+
config_attributes_v2=self._CONFIG_ATTRIBUTES_V2,
356+
config_attributes_v2_11=self._CONFIG_ATTRIBUTES_V2_11,
357+
include_extra=self.SUPPORTS_CONFIG_EXTRA,
358+
)
376359

377360
if self._has_lookaround_pattern():
378361
config_parameters["regex_engine"] = '"python-re"'
379362

380-
if isinstance(self.extra_template_data.get("config"), dict):
381-
for key, value in self.extra_template_data["config"].items():
382-
config_parameters[key] = value # noqa: PERF403
363+
if isinstance(config := self.extra_template_data.get("config"), dict):
364+
for key, value in config.items():
365+
config_parameters[key] = value
383366

384367
# Handle json_schema_extra from schema extensions (x-* fields)
385368
model_extras = self.extra_template_data.get("model_extras")
386369
if model_extras:
387-
existing = config_parameters.get("json_schema_extra") or {}
370+
existing = cast("dict[str, Any]", config_parameters.get("json_schema_extra") or {})
388371
config_parameters["json_schema_extra"] = {**existing, **model_extras}
389372

390373
if config_parameters:
@@ -395,40 +378,6 @@ def __init__( # noqa: PLR0913
395378

396379
self._process_validators()
397380

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-
432381
def _has_lookaround_pattern(self) -> bool:
433382
"""Check if any field has a regex pattern with lookaround assertions."""
434383
lookaround_regex = re.compile(r"\(\?<?[=!]")

‎src/datamodel_code_generator/model/pydantic_v2/dataclass.py‎

Lines changed: 8 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,12 @@
77

88
from typing import TYPE_CHECKING, Any, ClassVar
99

10-
from datamodel_code_generator.enums import TargetPydanticVersion
1110
from datamodel_code_generator.model import DataModel, DataModelFieldBase, _rebuild_model_with_datamodel_namespace
1211
from datamodel_code_generator.model.base import UNDEFINED
1312
from datamodel_code_generator.model.dataclass import _DataclassReuseMixin, has_field_assignment
14-
from datamodel_code_generator.model.pydantic_v2.base_model import (
13+
from datamodel_code_generator.model.pydantic_v2._config import (
1514
ConfigAttribute,
15+
build_base_config_parameters,
1616
)
1717
from datamodel_code_generator.model.pydantic_v2.base_model import (
1818
Constraints as _Constraints,
@@ -106,59 +106,17 @@ def __init__( # noqa: PLR0913
106106
self.dataclass_arguments["kw_only"] = True
107107
self._set_deprecated_decorator()
108108

109-
config_parameters: dict[str, Any] = {}
110-
111-
extra = self._get_config_extra()
112-
if extra:
113-
config_parameters["extra"] = extra
114-
115-
config_attributes = self._get_config_attributes()
116-
for from_, to, invert in config_attributes:
117-
if from_ in self.extra_template_data:
118-
config_parameters[to] = (
119-
not self.extra_template_data[from_] if invert else self.extra_template_data[from_]
120-
)
121-
122-
for data_type in self.all_data_types:
123-
if data_type.is_custom_type: # pragma: no cover
124-
config_parameters["arbitrary_types_allowed"] = True
125-
break
109+
config_parameters = build_base_config_parameters(
110+
extra_template_data=self.extra_template_data,
111+
all_data_types=self.all_data_types,
112+
config_attributes_v2=self._CONFIG_ATTRIBUTES_V2,
113+
config_attributes_v2_11=self._CONFIG_ATTRIBUTES_V2_11,
114+
)
126115

127116
if config_parameters:
128117
self._additional_imports.append(IMPORT_CONFIG_DICT)
129118
self.extra_template_data["config"] = config_parameters
130119

131-
def _get_config_attributes(self) -> list[ConfigAttribute]:
132-
"""Get config attributes based on target Pydantic version."""
133-
target_version = self.extra_template_data.get("target_pydantic_version")
134-
if target_version == TargetPydanticVersion.V2_11:
135-
return self._CONFIG_ATTRIBUTES_V2_11
136-
return self._CONFIG_ATTRIBUTES_V2
137-
138-
def _get_config_extra(self) -> str | None:
139-
"""Get extra field configuration for ConfigDict."""
140-
additional_properties = self.extra_template_data.get("additionalProperties")
141-
unevaluated_properties = self.extra_template_data.get("unevaluatedProperties")
142-
allow_extra_fields = self.extra_template_data.get("allow_extra_fields")
143-
extra_fields = self.extra_template_data.get("extra_fields")
144-
145-
config_extra = None
146-
if allow_extra_fields or extra_fields == "allow":
147-
config_extra = "'allow'"
148-
elif extra_fields == "forbid":
149-
config_extra = "'forbid'"
150-
elif extra_fields == "ignore":
151-
config_extra = "'ignore'"
152-
elif additional_properties is True:
153-
config_extra = "'allow'"
154-
elif additional_properties is False:
155-
config_extra = "'forbid'"
156-
elif unevaluated_properties is True:
157-
config_extra = "'allow'"
158-
elif unevaluated_properties is False:
159-
config_extra = "'forbid'"
160-
return config_extra
161-
162120

163121
if PYDANTIC_V2_DATACLASS_ALIAS_NEEDS_FALLBACK:
164122
import keyword

0 commit comments

Comments
 (0)