Skip to content

Commit 8a0d375

Browse files
koxudaxipre-commit-ci[bot]
koxudaxi
and
pre-commit-ci[bot]
authored
Fix allOf with $ref to root model losing constraints (#2676)
* feat: implement handling of allOf root model with constraints and merge mode * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: handle special path markers in allOf root model reference processing * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * feat: add new data types and constraints for allOf root model processing * feat: add new schema definitions and constraints for allOf processing * feat: add FormattedStringDatatype and ConflictingFormatAllOf to enhance allOf processing --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 599ae84 commit 8a0d375

6 files changed

Lines changed: 867 additions & 2 deletions

File tree

‎src/datamodel_code_generator/parser/jsonschema.py‎

Lines changed: 126 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@
6060
get_special_path,
6161
title_to_class_name,
6262
)
63-
from datamodel_code_generator.reference import ModelType, Reference, is_url
63+
from datamodel_code_generator.reference import SPECIAL_PATH_MARKER, ModelType, Reference, is_url
6464
from datamodel_code_generator.types import (
6565
ANY,
6666
DataType,
@@ -1162,6 +1162,43 @@ def _merge_primitive_schemas(self, items: list[JsonSchemaObject]) -> JsonSchemaO
11621162

11631163
return self.SCHEMA_OBJECT_TYPE.parse_obj(base_dict)
11641164

1165+
def _merge_primitive_schemas_for_allof(self, items: list[JsonSchemaObject]) -> JsonSchemaObject | None:
1166+
"""Merge primitive schemas for allOf, respecting allof_merge_mode setting."""
1167+
if len(items) == 1:
1168+
return items[0] # pragma: no cover
1169+
1170+
formats = {item.format for item in items if item.format}
1171+
if len(formats) > 1:
1172+
return None
1173+
1174+
merged_format = formats.pop() if formats else None
1175+
1176+
if self.allof_merge_mode != AllOfMergeMode.NoMerge:
1177+
merged = self._merge_primitive_schemas(items)
1178+
merged_dict = merged.dict(exclude_unset=True, by_alias=True)
1179+
if merged_format:
1180+
merged_dict["format"] = merged_format
1181+
return self.SCHEMA_OBJECT_TYPE.parse_obj(merged_dict)
1182+
1183+
base_dict: dict[str, Any] = {}
1184+
for item in items:
1185+
if item.type:
1186+
base_dict = item.dict(exclude_unset=True, by_alias=True)
1187+
break
1188+
1189+
for item in items:
1190+
for constraint_field in JsonSchemaObject.__constraint_fields__:
1191+
value = getattr(item, constraint_field, None)
1192+
if value is None:
1193+
value = item.extras.get(constraint_field)
1194+
if value is not None:
1195+
base_dict[constraint_field] = value
1196+
1197+
if merged_format:
1198+
base_dict["format"] = merged_format
1199+
1200+
return self.SCHEMA_OBJECT_TYPE.parse_obj(base_dict)
1201+
11651202
@staticmethod
11661203
def _intersect_constraint(field: str, val1: Any, val2: Any) -> Any: # noqa: PLR0911
11671204
"""Compute the intersection of two constraint values."""
@@ -1460,6 +1497,89 @@ def _schema_signature(self, prop_schema: JsonSchemaObject | bool) -> str | bool:
14601497
return prop_schema
14611498
return json.dumps(prop_schema.dict(exclude_unset=True, by_alias=True), sort_keys=True, default=repr)
14621499

1500+
def _is_root_model_schema(self, obj: JsonSchemaObject) -> bool: # noqa: PLR6301
1501+
"""Check if schema represents a root model (primitive type with constraints).
1502+
1503+
Based on parse_raw_obj() else branch conditions. Returns True when
1504+
the schema would be processed by parse_root_type().
1505+
"""
1506+
if obj.is_array:
1507+
return False
1508+
if obj.allOf or obj.oneOf or obj.anyOf:
1509+
return False
1510+
if obj.properties:
1511+
return False
1512+
if obj.patternProperties:
1513+
return False
1514+
if obj.type == "object":
1515+
return False
1516+
return not obj.enum
1517+
1518+
def _handle_allof_root_model_with_constraints( # noqa: PLR0911, PLR0912
1519+
self,
1520+
name: str,
1521+
obj: JsonSchemaObject,
1522+
path: list[str],
1523+
) -> DataType | None:
1524+
"""Handle allOf that combines a root model $ref with additional constraints.
1525+
1526+
This handler is for generating a root model from a root model reference.
1527+
Object inheritance (with properties) is handled by existing _parse_all_of_item() path.
1528+
Only applies to named schema definitions, not inline properties.
1529+
"""
1530+
for path_element in path:
1531+
if SPECIAL_PATH_MARKER in path_element:
1532+
return None # pragma: no cover
1533+
1534+
ref_items = [item for item in obj.allOf if item.ref]
1535+
1536+
if len(ref_items) != 1:
1537+
return None
1538+
1539+
ref_item = ref_items[0]
1540+
ref_value = ref_item.ref
1541+
if ref_value is None:
1542+
return None # pragma: no cover
1543+
1544+
if ref_item.has_ref_with_schema_keywords:
1545+
ref_schema = self._merge_ref_with_schema(ref_item)
1546+
else:
1547+
ref_schema = self._load_ref_schema_object(ref_value)
1548+
1549+
if not self._is_root_model_schema(ref_schema):
1550+
return None
1551+
1552+
constraint_items: list[JsonSchemaObject] = []
1553+
for item in obj.allOf:
1554+
if item.ref:
1555+
continue
1556+
if item.properties or item.items:
1557+
return None
1558+
if item.has_constraint or item.type or item.format:
1559+
if item.type and ref_schema.type:
1560+
compatible_type_pairs = {
1561+
("integer", "number"),
1562+
("number", "integer"),
1563+
}
1564+
if item.type != ref_schema.type and (item.type, ref_schema.type) not in compatible_type_pairs:
1565+
return None
1566+
constraint_items.append(item)
1567+
1568+
if not constraint_items:
1569+
return None
1570+
1571+
all_items = [ref_schema, *constraint_items]
1572+
merged_schema = self._merge_primitive_schemas_for_allof(all_items)
1573+
if merged_schema is None:
1574+
return None
1575+
1576+
if obj.description:
1577+
merged_dict = merged_schema.dict(exclude_unset=True, by_alias=True)
1578+
merged_dict["description"] = obj.description
1579+
merged_schema = self.SCHEMA_OBJECT_TYPE.parse_obj(merged_dict)
1580+
1581+
return self.parse_root_type(name, merged_schema, path)
1582+
14631583
def _merge_all_of_object(self, obj: JsonSchemaObject) -> JsonSchemaObject | None:
14641584
"""Merge allOf items when they share object properties to avoid duplicate models.
14651585
@@ -1869,6 +1989,11 @@ def parse_all_of(
18691989
base_classes=[],
18701990
required=[],
18711991
)
1992+
1993+
root_model_result = self._handle_allof_root_model_with_constraints(name, obj, path)
1994+
if root_model_result is not None:
1995+
return root_model_result
1996+
18721997
fields: list[DataModelFieldBase] = []
18731998
base_classes: list[Reference] = []
18741999
required: list[str] = []

‎tests/data/expected/main/jsonschema/allof_root_model_constraints.py‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
# generated by datamodel-codegen:
2+
# filename: allof_root_model_constraints.json
3+
# timestamp: 2019-07-26T00:00:00+00:00
4+
5+
from __future__ import annotations
6+
7+
from typing import Any, Dict, List, Optional
8+
9+
from pydantic import BaseModel, EmailStr, Field, conint, constr
10+
11+
12+
class StringDatatype(BaseModel):
13+
__root__: constr(regex=r'^\S(.*\S)?$') = Field(
14+
..., description='A base string type.'
15+
)
16+
17+
18+
class ConstrainedStringDatatype(BaseModel):
19+
__root__: constr(regex=r'^[A-Z].*', min_length=1) = Field(
20+
..., description='A constrained string.'
21+
)
22+
23+
24+
class IntegerDatatype(BaseModel):
25+
__root__: int = Field(..., description='A whole number.')
26+
27+
28+
class NonNegativeIntegerDatatype(BaseModel):
29+
__root__: conint(ge=0) = Field(..., description='Non-negative integer.')
30+
31+
32+
class BoundedIntegerDatatype(BaseModel):
33+
__root__: conint(ge=0, le=100) = Field(
34+
..., description='Integer between 0 and 100.'
35+
)
36+
37+
38+
class EmailDatatype(BaseModel):
39+
__root__: EmailStr = Field(..., description='Email with format.')
40+
41+
42+
class FormattedStringDatatype(BaseModel):
43+
__root__: EmailStr = Field(..., description='A string with email format.')
44+
45+
46+
class ObjectBase(BaseModel):
47+
id: Optional[int] = None
48+
49+
50+
class ObjectWithAllOf(ObjectBase):
51+
name: Optional[str] = None
52+
53+
54+
class MultiRefAllOf(BaseModel):
55+
pass
56+
57+
58+
class NoConstraintAllOf(BaseModel):
59+
pass
60+
61+
62+
class IncompatibleTypeAllOf(BaseModel):
63+
pass
64+
65+
66+
class ConstraintWithProperties(BaseModel):
67+
extra: Optional[str] = None
68+
69+
70+
class ConstraintWithItems(BaseModel):
71+
pass
72+
73+
74+
class NumberIntegerCompatible(BaseModel):
75+
__root__: conint(ge=0) = Field(
76+
..., description='Number and integer are compatible.'
77+
)
78+
79+
80+
class RefWithSchemaKeywords(BaseModel):
81+
__root__: constr(regex=r'^\S(.*\S)?$', min_length=5, max_length=100) = Field(
82+
..., description='Ref with additional schema keywords.'
83+
)
84+
85+
86+
class ArrayDatatype(BaseModel):
87+
__root__: List[str]
88+
89+
90+
class RefToArrayAllOf(BaseModel):
91+
pass
92+
93+
94+
class ObjectNoPropsDatatype(BaseModel):
95+
pass
96+
97+
98+
class RefToObjectNoPropsAllOf(ObjectNoPropsDatatype):
99+
pass
100+
101+
102+
class PatternPropsDatatype(BaseModel):
103+
__root__: Dict[constr(regex=r'^S_'), str]
104+
105+
106+
class RefToPatternPropsAllOf(BaseModel):
107+
pass
108+
109+
110+
class NestedAllOfDatatype(BaseModel):
111+
pass
112+
113+
114+
class RefToNestedAllOfAllOf(NestedAllOfDatatype):
115+
pass
116+
117+
118+
class ConstraintsOnlyDatatype(BaseModel):
119+
__root__: Any = Field(..., description='Constraints only, no type.')
120+
121+
122+
class RefToConstraintsOnlyAllOf(BaseModel):
123+
__root__: Any = Field(..., description='Ref to constraints-only schema.')
124+
125+
126+
class NoDescriptionAllOf(BaseModel):
127+
__root__: constr(regex=r'^\S(.*\S)?$', min_length=5) = Field(
128+
..., description='A base string type.'
129+
)
130+
131+
132+
class EmptyConstraintItemAllOf(BaseModel):
133+
__root__: constr(regex=r'^\S(.*\S)?$', max_length=50) = Field(
134+
..., description='AllOf with empty constraint item.'
135+
)
136+
137+
138+
class ConflictingFormatAllOf(BaseModel):
139+
pass
140+
141+
142+
class Model(BaseModel):
143+
name: Optional[ConstrainedStringDatatype] = None
144+
count: Optional[NonNegativeIntegerDatatype] = None
145+
percentage: Optional[BoundedIntegerDatatype] = None
146+
email: Optional[EmailDatatype] = None
147+
obj: Optional[ObjectWithAllOf] = None
148+
multi: Optional[MultiRefAllOf] = None
149+
noconstraint: Optional[NoConstraintAllOf] = None
150+
incompatible: Optional[IncompatibleTypeAllOf] = None
151+
withprops: Optional[ConstraintWithProperties] = None
152+
withitems: Optional[ConstraintWithItems] = None
153+
numint: Optional[NumberIntegerCompatible] = None
154+
refwithkw: Optional[RefWithSchemaKeywords] = None
155+
refarr: Optional[RefToArrayAllOf] = None
156+
refobjnoprops: Optional[RefToObjectNoPropsAllOf] = None
157+
refpatternprops: Optional[RefToPatternPropsAllOf] = None
158+
refnestedallof: Optional[RefToNestedAllOfAllOf] = None
159+
refconstraintsonly: Optional[RefToConstraintsOnlyAllOf] = None
160+
nodescription: Optional[NoDescriptionAllOf] = None
161+
emptyconstraint: Optional[EmptyConstraintItemAllOf] = None
162+
conflictingformat: Optional[ConflictingFormatAllOf] = None

0 commit comments

Comments
 (0)