fix: enforce array contains/minContains/maxContains on generated models (#57)

* fix: enforce array contains/minContains/maxContains on generated models

datamodel-code-generator drops contains/minContains/maxContains, so the
generated Totals (a bare list[Total] alias) accepts arrays that violate
totals.json's "exactly one subtotal and one total" rule. Snapshot the pristine
schemas before preprocessing (which merges allOf and would drop the second
contains), derive every predicate from contains.properties.*.const, and thread
a pydantic AfterValidator enforcing all bounds into the alias metadata — base
model and request variants alike. Also fix the moved description import so the
SDK-guarded codegen tests run. Data-driven, idempotent, dependency-free.

* chore: ignore N801 in generated schemas in ruff config

---------

Co-authored-by: Vishal Katyal <vishal@katyal.ai>
Co-authored-by: damaz91 <federico.damato91@gmail.com>

Author: 3 people

generate_models.sh

@@ -46,6 +46,15 @@ OUTPUT_DIR="src/ucp_sdk/models/schemas"
 # Schema directory (relative to this script)
 SCHEMA_DIR="ucp/source/schemas"
 
+# Snapshot the pristine schemas before preprocessing. postprocess_models.py
+# reads array contains/minContains/maxContains from these originals because
+# preprocessing merges allOf branches and a JSON node holds only one contains,
+# so a second contains keyword (e.g. "exactly one total") would otherwise be
+# silently dropped before the post-processor could see it.
+RAW_SCHEMA_DIR="ucp/raw_schemas"
+rm -rf "$RAW_SCHEMA_DIR"
+cp -R "$SCHEMA_DIR" "$RAW_SCHEMA_DIR"
+
 echo "Preprocessing schemas..."
 uv run python preprocess_schemas.py
 

postprocess_models.py

@@ -27,6 +27,27 @@
 keys (``model_extra``) — an explicit null is a present key, and unknown keys
 on ``extra="allow"`` models count too.
 
+``contains`` / ``minContains`` / ``maxContains`` on an array schema is likewise
+dropped by the generator: ``totals.json`` requires *exactly one* ``subtotal``
+*and exactly one* ``total`` entry, but the generated ``Totals`` is a bare
+``list[Total]`` alias, so an empty array (or one missing either required entry,
+or with duplicates) validates in violation of the schema. An array root is
+emitted as a ``TypeAliasType`` wrapping ``Annotated[list[...], ...]`` rather
+than a ``BaseModel`` subclass, so ``model_validator`` cannot apply; this script
+instead injects a module-level counting function, threaded into the alias
+metadata as a ``pydantic.AfterValidator``. Every predicate is derived from
+``contains.properties.<field>.const`` — nothing is hard-coded — and one function
+enforces *all* of a schema's contains bounds.
+
+The pristine (pre-preprocessing) schemas are read for this: ``totals.json``
+carries its two containment rules as two ``allOf`` branches, and
+``preprocess_schemas.py`` merges ``allOf`` into the root, where a JSON node can
+hold only one ``contains`` — so the second (``total``) would be lost if the
+preprocessed output were scanned. generate_models.sh snapshots the originals to
+``ucp/raw_schemas`` before preprocessing for exactly this reason. The bound is
+applied to the base model and to its generated request variants (linked by file
+stem), and travels wherever the alias is reused as a field type.
+
 Runs from generate_models.sh between generation and formatting; idempotent.
 """
 
@@ -36,6 +57,10 @@
 from pathlib import Path
 
 SCHEMA_DIR = Path("ucp/source/schemas")
+# Pristine schemas snapshotted by generate_models.sh before preprocessing.
+# Array contains bounds are read from here, not SCHEMA_DIR, because
+# preprocessing merges allOf and can drop a second contains keyword.
+RAW_SCHEMA_DIR = Path("ucp/raw_schemas")
 OUTPUT_DIR = Path("src/ucp_sdk/models/schemas")
 
 _MARKER = "_enforce_min_properties"
@@ -79,13 +104,15 @@ def find_root_min_properties(schema_dir):
     return found
 
 
-def _ensure_validator_import(source):
-    """Add model_validator to the existing pydantic import if missing."""
-    if re.search(r"^from pydantic import .*\bmodel_validator\b", source, re.M):
+def _ensure_pydantic_import(source, symbol):
+    """Add ``symbol`` to the ``from pydantic import`` line if absent."""
+    if re.search(
+        rf"^from pydantic import .*\b{re.escape(symbol)}\b", source, re.M
+    ):
         return source
     return re.sub(
         r"^(from pydantic import [^\n]+)$",
-        lambda m: f"{m.group(1)}, model_validator",
+        lambda m: f"{m.group(1)}, {symbol}",
         source,
         count=1,
         flags=re.M,
@@ -112,17 +139,196 @@ def inject_min_properties(source, class_name, minimum):
     body = source[:end].rstrip("\n")
     rest = source[end:]
     out = body + "\n" + method + ("\n" + rest if rest else "")
-    return _ensure_validator_import(out)
+    return _ensure_pydantic_import(out, "model_validator")
 
 
-def main():
-    """Main entry point to scan schemas and patch generated models."""
+def _extract_contains_groups(schema, path=None):
+    """Collect every array ``contains`` group from a schema's root + allOf.
+
+    Each group is ``{"pairs": [(field, const), ...], "min": int,
+    "max": int | None}``, derived from ``contains.properties.<field>.const``
+    with its ``minContains`` / ``maxContains`` bounds. A ``contains`` keyword
+    may sit at the schema root or inside any ``allOf`` branch; each contributes
+    a group, so "exactly one subtotal and one total" yields two. The predicate
+    is read from the schema, never hard-coded.
+    """
+    nodes = [schema]
+    if isinstance(schema.get("allOf"), list):
+        nodes.extend(n for n in schema["allOf"] if isinstance(n, dict))
+    groups = []
+    for node in nodes:
+        contains = node.get("contains")
+        if not isinstance(contains, dict):
+            continue
+        props = contains.get("properties")
+        pairs = []
+        if isinstance(props, dict):
+            for field, spec in props.items():
+                if isinstance(spec, dict) and "const" in spec:
+                    pairs.append((field, spec["const"]))
+        if not pairs:
+            if path is not None:
+                sys.stderr.write(
+                    f"  ! {path}: contains predicate has no "
+                    "properties.*.const; cannot derive a check\n"
+                )
+            continue
+        # JSON Schema: minContains defaults to 1 when contains is present.
+        groups.append(
+            {
+                "pairs": pairs,
+                "min": node.get("minContains", 1),
+                "max": node.get("maxContains"),
+            }
+        )
+    return groups
+
+
+def find_array_contains_constraints(schema_dir):
+    """Map file stem -> ``{"title": str, "groups": [...]}`` for array schemas.
+
+    Keyed by file stem (not title) so a base schema can be linked to its
+    generated request variants, whose stems extend it (``totals`` ->
+    ``totals_create_request``). Scanned against the *pristine* schemas
+    (``RAW_SCHEMA_DIR``); see the module docstring for why the preprocessed
+    output must not be used here.
+    """
+    found = {}
+    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
+        # ``contains`` only constrains arrays; skip anything else.
+        if schema.get("type") != "array" and "items" not in schema:
+            continue
+        groups = _extract_contains_groups(schema, path)
+        if not groups:
+            continue
+        title = schema.get("title")
+        if not title:
+            sys.stderr.write(
+                f"  ! {path}: array contains constraint but no title; "
+                "cannot map to a model\n"
+            )
+            continue
+        found[path.stem] = {"title": title, "groups": groups}
+    return found
+
+
+def _alias_name(title):
+    """Derive the generated alias name from a schema title (drop spaces)."""
+    return "".join(title.split())
+
+
+def _snake_name(name):
+    """CamelCase alias -> snake_case suffix for a unique function name."""
+    return re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
+
+
+def _predicate_expr(pairs):
+    """Build a per-item boolean expression matching all (field, const) pairs.
+
+    Items are ``Total`` instances after inner validation, but a mapping is
+    handled too so the check is robust regardless of the item representation.
+    """
+    parts = []
+    for field, const in pairs:
+        parts.append(
+            f"(_item.get({field!r}) if isinstance(_item, dict) "
+            f"else getattr(_item, {field!r}, None)) == {const!r}"
+        )
+    return " and ".join(parts)
+
+
+def _build_contains_function(func_name, groups):
+    """Render the module-level ``AfterValidator`` counting function."""
+    lines = [
+        f"def {func_name}(value):",
+        '    """JSON Schema contains/minContains/maxContains (see #49)."""',
+    ]
+    for index, group in enumerate(groups):
+        count = "_matched" if len(groups) == 1 else f"_matched_{index}"
+        desc = ", ".join(f"{f}=={c!r}" for f, c in group["pairs"])
+        lines += [
+            f"    {count} = sum(",
+            "        1",
+            "        for _item in value",
+            f"        if {_predicate_expr(group['pairs'])}",
+            "    )",
+        ]
+        minimum = group["min"]
+        noun = "entry" if minimum == 1 else "entries"
+        lines += [
+            f"    if {count} < {minimum}:",
+            "        raise ValueError(",
+            f'            "Array must contain at least {minimum} {noun} "',
+            f'            "matching {desc} (schema minContains={minimum})"',
+            "        )",
+        ]
+        maximum = group["max"]
+        if maximum is not None:
+            noun = "entry" if maximum == 1 else "entries"
+            lines += [
+                f"    if {count} > {maximum}:",
+                "        raise ValueError(",
+                f'            "Array must contain at most {maximum} {noun} "',
+                f'            "matching {desc} (schema maxContains={maximum})"',
+                "        )",
+            ]
+    lines.append("    return value")
+    return "\n".join(lines) + "\n"
+
+
+def inject_array_contains(source, alias_name, groups):
+    """Thread an ``AfterValidator`` into ``alias_name``'s alias metadata.
+
+    Array roots are emitted as ``NAME = TypeAliasType("NAME", Annotated[...])``,
+    not a ``BaseModel`` subclass, so the constraint is enforced by inserting
+    ``AfterValidator(<fn>)`` into the ``Annotated[...]`` metadata and defining
+    ``<fn>`` just above the assignment. Idempotent via the function name.
+    """
+    func_name = f"_enforce_contains_{_snake_name(alias_name)}"
+    if f"def {func_name}(" in source:
+        return source
+    assign_re = re.compile(rf"^{re.escape(alias_name)} = TypeAliasType\(", re.M)
+    match = assign_re.search(source)
+    if not match:
+        return source
+    ann_start = source.find("Annotated[", match.end())
+    if ann_start == -1:
+        return source
+    # Bracket-match to the ``]`` that closes ``Annotated[``.
+    depth = 0
+    close = None
+    for pos in range(ann_start + len("Annotated"), len(source)):
+        char = source[pos]
+        if char == "[":
+            depth += 1
+        elif char == "]":
+            depth -= 1
+            if depth == 0:
+                close = pos
+                break
+    if close is None:
+        return source
+    out = source[:close] + f", AfterValidator({func_name})" + source[close:]
+    func_src = _build_contains_function(func_name, groups)
+    insert_at = assign_re.search(out).start()
+    out = out[:insert_at] + func_src + "\n\n" + out[insert_at:]
+    return _ensure_pydantic_import(out, "AfterValidator")
+
+
+def _patch_min_properties():
+    """Inject minProperties validators; return (patched_count, exit_code)."""
     constraints = find_root_min_properties(SCHEMA_DIR)
     if not constraints:
         sys.stdout.write(
             "postprocess: no root-level minProperties constraints found\n"
         )
-        return 0
+        return 0, 0
     patched = 0
     for title, minimum in sorted(constraints.items()):
         hits = []
@@ -142,9 +348,90 @@ def main():
                 f"  ! '{title}' has no generated class; "
                 "constraint not enforced\n"
             )
-            return 1
-    sys.stdout.write(f"postprocess: {patched} module(s) patched\n")
-    return 0
+            return patched, 1
+    return patched, 0
+
+
+def _array_contains_targets():
+    """Resolve ``title -> groups`` for every model needing a contains bound.
+
+    The authoritative (complete) groups come from the pristine schemas. The
+    preprocessed tree is consulted only to enumerate which models actually
+    exist — the base plus its generated request variants — so each variant
+    inherits its base schema's full set of containment rules. Variants are
+    linked to their base by file stem (``totals_create_request`` -> ``totals``).
+    """
+    raw = find_array_contains_constraints(RAW_SCHEMA_DIR)
+    if not raw:
+        # Fallback keeps a standalone run working if the snapshot is absent,
+        # though the pipeline always provides it (see module docstring).
+        raw = find_array_contains_constraints(SCHEMA_DIR)
+        if raw:
+            sys.stderr.write(
+                f"  ! {RAW_SCHEMA_DIR} missing; falling back to preprocessed "
+                "schemas (multi-branch contains may be incomplete)\n"
+            )
+    if not raw:
+        return {}
+    raw_stems = sorted(raw, key=len, reverse=True)
+    targets = {}
+    # Enumerate base + variants from the preprocessed tree; attach raw groups.
+    for stem, info in find_array_contains_constraints(SCHEMA_DIR).items():
+        origin = next(
+            (s for s in raw_stems if stem == s or stem.startswith(s + "_")),
+            None,
+        )
+        if origin is not None:
+            targets[info["title"]] = raw[origin]["groups"]
+    # Defensive: cover each raw base title even if the preprocessed base lost
+    # its contains entirely.
+    for info in raw.values():
+        targets.setdefault(info["title"], info["groups"])
+    return targets
+
+
+def _patch_array_contains():
+    """Inject array-contains validators; return (patched_count, exit_code)."""
+    targets = _array_contains_targets()
+    if not targets:
+        sys.stdout.write("postprocess: no array contains constraints found\n")
+        return 0, 0
+    patched = 0
+    for title, groups in sorted(targets.items()):
+        alias = _alias_name(title)
+        hits = []
+        for path in sorted(OUTPUT_DIR.rglob("*.py")):
+            source = path.read_text(encoding="utf-8")
+            if not re.search(
+                rf"^{re.escape(alias)} = TypeAliasType\(", source, re.M
+            ):
+                continue
+            updated = inject_array_contains(source, alias, groups)
+            if updated != source:
+                path.write_text(updated, encoding="utf-8")
+                patched += 1
+            hits.append(path)
+        preds = "; ".join(
+            " & ".join(f"{f}=={c!r}" for f, c in g["pairs"]) for g in groups
+        )
+        label = ", ".join(str(h) for h in hits) or "NO GENERATED ALIAS FOUND"
+        sys.stdout.write(f"  contains [{preds}] on '{title}' -> {label}\n")
+        if not hits:
+            sys.stderr.write(
+                f"  ! '{title}' has no generated alias; "
+                "constraint not enforced\n"
+            )
+            return patched, 1
+    return patched, 0
+
+
+def main():
+    """Main entry point to scan schemas and patch generated models."""
+    patched_mp, rc_mp = _patch_min_properties()
+    patched_ac, rc_ac = _patch_array_contains()
+    total = patched_mp + patched_ac
+    sys.stdout.write(f"postprocess: {total} module(s) patched\n")
+    return rc_mp or rc_ac
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -69,7 +69,7 @@ skip-magic-trailing-comma = false
 line-ending = "auto"
 
 [tool.ruff.lint.per-file-ignores]
-"src/ucp_sdk/models/schemas/**/*.py" = ["E501", "D"]
+"src/ucp_sdk/models/schemas/**/*.py" = ["E501", "D", "N801"]
 
 [tool.ruff.lint]
 select = ["E", "F", "W", "B", "C4", "SIM", "N", "UP", "D", "PTH", "T20"]