fix: enforce conditional required fields in generated models (#69)

Author: FanouZeng-TT

postprocess_models.py

@@ -67,6 +67,13 @@
   injects a ``field_validator(mode="after")`` into each generated class that
   declares a matching list field.
 
+* Simple conditional ``required`` constraints are dropped: pagination requires
+  ``cursor`` when ``has_next_page`` is true, but the generated response model
+  always treats it as optional. The script accepts only an unambiguous single
+  required discriminator using ``const``/``enum`` and a ``then.required`` list,
+  then injects a ``model_validator(mode="after")``. More complex conditions are
+  skipped rather than approximated.
+
 Runs from generate_models.sh between generation and formatting; idempotent.
 """
 
@@ -117,6 +124,24 @@ def {marker}(self):
 
 _UNIQUE_MARKER = "_enforce_unique_items"
 
+_CONDITIONAL_REQUIRED_MARKER = "_enforce_conditional_required"
+
+_CONDITIONAL_REQUIRED_TEMPLATE = '''
+    @model_validator(mode="after")
+    def {marker}(self):
+        """JSON Schema if/then: enforce conditionally required fields."""
+        rules = {rules!r}
+        for rule in rules:
+            if getattr(self, rule["discriminator"], None) not in rule["values"]:
+                continue
+            for field in rule["required"]:
+                if field not in self.model_fields_set:
+                    raise ValueError(
+                        f"Field {{field!r}} is required by a schema condition"
+                    )
+        return self
+'''
+
 _UNIQUE_VALIDATOR_TEMPLATE = '''
     @field_validator("{field}", mode="after")
     def {marker}_{field}(cls, value):  # noqa: N805
@@ -528,6 +553,137 @@ def inject_array_contains(source, alias_name, groups):
     return _ensure_pydantic_import(out, "AfterValidator")
 
 
+def find_conditional_required(schema_dir):
+    """Map generated class names to simple if/then required rules."""
+    rules_by_class = {}
+
+    def describe(node, properties):
+        if not isinstance(node, dict) or set(node) != {"if", "then"}:
+            return None
+        condition = node["if"]
+        consequence = node["then"]
+        if (
+            not isinstance(condition, dict)
+            or set(condition) != {"properties", "required"}
+            or not isinstance(consequence, dict)
+            or set(consequence) != {"required"}
+        ):
+            return None
+        condition_props = condition["properties"]
+        condition_required = condition["required"]
+        consequence_required = consequence["required"]
+        if (
+            not isinstance(condition_props, dict)
+            or len(condition_props) != 1
+            or not isinstance(condition_required, list)
+            or len(condition_required) != 1
+            or not isinstance(consequence_required, list)
+            or not consequence_required
+        ):
+            return None
+        discriminator, predicate = next(iter(condition_props.items()))
+        if condition_required != [discriminator] or not isinstance(
+            predicate, dict
+        ):
+            return None
+        if set(predicate) == {"const"}:
+            values = [predicate["const"]]
+        elif (
+            set(predicate) == {"enum"}
+            and isinstance(predicate["enum"], list)
+            and predicate["enum"]
+        ):
+            values = predicate["enum"]
+        else:
+            return None
+        if (
+            discriminator not in properties
+            or any(
+                not isinstance(name, str) or name not in properties
+                for name in consequence_required
+            )
+            or any(
+                not isinstance(value, (str, int, float, bool))
+                for value in values
+            )
+        ):
+            return None
+        return {
+            "discriminator": discriminator,
+            "values": values,
+            "required": sorted(consequence_required),
+        }
+
+    def walk(node, current_class_name, path_str):
+        if not isinstance(node, dict):
+            return
+        if isinstance(node.get("title"), str):
+            current_class_name = _alias_name(node["title"])
+        properties = node.get("properties")
+        then = node.get("then")
+        is_required_rule = isinstance(then, dict) and "required" in then
+        if isinstance(properties, dict) and is_required_rule:
+            if "else" in node:
+                rule = None
+            else:
+                rule = describe(
+                    {key: node[key] for key in ("if", "then") if key in node},
+                    properties,
+                )
+            if rule is None:
+                sys.stderr.write(
+                    f"  ! {path_str}: unsupported conditional required rule; skipped\n"
+                )
+            elif current_class_name is not None:
+                rules_by_class.setdefault(current_class_name, []).append(rule)
+        if isinstance(properties, dict):
+            for name, prop in properties.items():
+                walk(prop, _to_camel_case(name), path_str)
+        defs = node.get("$defs")
+        if isinstance(defs, dict):
+            for def_name, def_node in defs.items():
+                walk(def_node, _to_camel_case(def_name), path_str)
+        for key in ("allOf", "anyOf", "oneOf"):
+            if isinstance(node.get(key), list):
+                for item in node[key]:
+                    walk(item, current_class_name, path_str)
+
+    for path in sorted(Path(schema_dir).rglob("*.json")):
+        try:
+            schema = json.loads(path.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            continue
+        if not isinstance(schema, dict):
+            continue
+        root_title = schema.get("title")
+        initial_class = (
+            _alias_name(root_title) if root_title else _to_camel_case(path.stem)
+        )
+        walk(schema, initial_class, str(path))
+    return rules_by_class
+
+
+def inject_conditional_required(source, class_name, rules):
+    """Inject simple conditional-required checks into one generated class."""
+    class_re = re.compile(rf"^class {re.escape(class_name)}\(", re.M)
+    match = class_re.search(source)
+    if not match:
+        return source
+    tail = re.compile(r"^\S", re.M)
+    end_match = tail.search(source, match.end())
+    end = end_match.start() if end_match else len(source)
+    if f"def {_CONDITIONAL_REQUIRED_MARKER}(" in source[match.start() : end]:
+        return source
+    method = _CONDITIONAL_REQUIRED_TEMPLATE.format(
+        marker=_CONDITIONAL_REQUIRED_MARKER,
+        rules=rules,
+    )
+    body = source[:end].rstrip("\n")
+    rest = source[end:]
+    out = body + "\n" + method + ("\n" + rest if rest else "")
+    return _ensure_pydantic_import(out, "model_validator")
+
+
 def find_unique_items_fields(schema_dir):
     """Map generated class names to fields carrying ``uniqueItems``.
 
@@ -787,6 +943,39 @@ def _patch_array_contains():
     return patched, 0
 
 
+def _patch_conditional_required():
+    """Inject conditional-required validators; return counts and status."""
+    rules_by_class = find_conditional_required(SCHEMA_DIR)
+    if not rules_by_class:
+        sys.stdout.write(
+            "postprocess: no simple conditional required rules found\n"
+        )
+        return 0, 0
+    patched = 0
+    for class_name, rules in sorted(rules_by_class.items()):
+        hits = []
+        for path in sorted(OUTPUT_DIR.rglob("*.py")):
+            source = path.read_text(encoding="utf-8")
+            if not re.search(
+                rf"^class {re.escape(class_name)}\(", source, re.M
+            ):
+                continue
+            updated = inject_conditional_required(source, class_name, rules)
+            if updated != source:
+                path.write_text(updated, encoding="utf-8")
+                patched += 1
+            hits.append(path)
+        label = (
+            ", ".join(str(path) for path in hits) or "NO GENERATED CLASS FOUND"
+        )
+        sys.stdout.write(
+            f"  conditional required on '{class_name}' -> {label}\n"
+        )
+        if not hits:
+            return patched, 1
+    return patched, 0
+
+
 def _patch_unique_items():
     """Inject uniqueItems validators; return (patched_count, exit_code)."""
     unique_fields_by_class = find_unique_items_fields(SCHEMA_DIR)
@@ -820,10 +1009,11 @@ def main():
     patched_mp, rc_mp = _patch_min_properties()
     patched_pn, rc_pn = _patch_property_names()
     patched_ac, rc_ac = _patch_array_contains()
+    patched_cr, rc_cr = _patch_conditional_required()
     patched_ui, rc_ui = _patch_unique_items()
-    total = patched_mp + patched_pn + patched_ac + patched_ui
+    total = patched_mp + patched_pn + patched_ac + patched_cr + patched_ui
     sys.stdout.write(f"postprocess: {total} module(s) patched\n")
-    return rc_mp or rc_pn or rc_ac or rc_ui
+    return rc_mp or rc_pn or rc_ac or rc_cr or rc_ui
 
 
 if __name__ == "__main__":

src/ucp_sdk/models/schemas/shopping/types/pagination.py

@@ -18,7 +18,7 @@
 
 from __future__ import annotations
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 
 class Pagination(BaseModel):
@@ -69,3 +69,23 @@ class Response(BaseModel):
     """
     Total number of matching items, if available.
     """
+
+    @model_validator(mode="after")
+    def _enforce_conditional_required(self):
+        """JSON Schema if/then: enforce conditionally required fields."""
+        rules = [
+            {
+                "discriminator": "has_next_page",
+                "values": [True],
+                "required": ["cursor"],
+            }
+        ]
+        for rule in rules:
+            if getattr(self, rule["discriminator"], None) not in rule["values"]:
+                continue
+            for field in rule["required"]:
+                if field not in self.model_fields_set:
+                    raise ValueError(
+                        f"Field {field!r} is required by a schema condition"
+                    )
+        return self

tests/test_codegen_pipeline.py

@@ -984,6 +984,95 @@ def test_injected_validator_enforces_pattern(self):
         signals_cls.model_validate({"com.example.ok": "v"})
 
 
+class ConditionalRequiredInjectorTest(unittest.TestCase):
+    """Simple JSON Schema if/then required constraints are restored."""
+
+    MODULE = (
+        "from __future__ import annotations\n"
+        "\n"
+        "from pydantic import BaseModel, ConfigDict\n"
+        "\n"
+        "\n"
+        "class Response(BaseModel):\n"
+        '    model_config = ConfigDict(extra="allow")\n'
+        "    cursor: str | None = None\n"
+        "    has_next_page: bool\n"
+    )
+    RULES = [
+        {
+            "discriminator": "has_next_page",
+            "values": [True],
+            "required": ["cursor"],
+        }
+    ]
+
+    def test_schema_scan_maps_nested_definition_to_generated_class(self):
+        schema = {
+            "title": "Pagination",
+            "type": "object",
+            "$defs": {
+                "response": {
+                    "type": "object",
+                    "properties": {
+                        "cursor": {"type": "string"},
+                        "has_next_page": {"type": "boolean"},
+                    },
+                    "if": {
+                        "properties": {"has_next_page": {"const": True}},
+                        "required": ["has_next_page"],
+                    },
+                    "then": {"required": ["cursor"]},
+                }
+            },
+        }
+        with tempfile.TemporaryDirectory() as tmp:
+            Path(tmp, "pagination.json").write_text(json.dumps(schema))
+            found = postprocess_models.find_conditional_required(Path(tmp))
+        self.assertEqual(found, {"Response": self.RULES})
+
+    def test_schema_scan_skips_else_branches(self):
+        schema = {
+            "title": "Response",
+            "type": "object",
+            "properties": {
+                "cursor": {"type": "string"},
+                "has_next_page": {"type": "boolean"},
+            },
+            "if": {
+                "properties": {"has_next_page": {"const": True}},
+                "required": ["has_next_page"],
+            },
+            "then": {"required": ["cursor"]},
+            "else": {"required": ["other"]},
+        }
+        with tempfile.TemporaryDirectory() as tmp:
+            Path(tmp, "response.json").write_text(json.dumps(schema))
+            found = postprocess_models.find_conditional_required(Path(tmp))
+        self.assertEqual(found, {})
+
+    def test_injection_is_idempotent(self):
+        once = postprocess_models.inject_conditional_required(
+            self.MODULE, "Response", self.RULES
+        )
+        twice = postprocess_models.inject_conditional_required(
+            once, "Response", self.RULES
+        )
+        self.assertEqual(once, twice)
+
+    @unittest.skipUnless(HAVE_SDK, "executing the module needs pydantic")
+    def test_injected_validator_enforces_conditional_required(self):
+        out = postprocess_models.inject_conditional_required(
+            self.MODULE, "Response", self.RULES
+        )
+        namespace: dict = {}
+        exec(compile(out, "<injected>", "exec"), namespace)  # noqa: S102
+        response = namespace["Response"]
+        with self.assertRaises(ValidationError):
+            response(has_next_page=True)
+        response(has_next_page=True, cursor="next-page")
+        response(has_next_page=False)
+
+
 class InjectorTest(unittest.TestCase):
     """The post-generation injector's own behavior."""