Skip to content

Commit 3228b23

Browse files
committed
feat: support capability extensions in codegen and cart-to-checkout conversion
- Generate request variants across all capability extension $defs in preprocess_schemas.py - Inject _enforce_cart_conversion validator in postprocess_models.py to allow cart_id or line_items in cart checkout creation - Fix trailing comma issue when injecting AfterValidator into multiline Annotated aliases - Add CapabilityExtensionSemanticTest to test_codegen_pipeline.py - Regenerate models against UCP version 2026-04-08
1 parent d790602 commit 3228b23

55 files changed

Lines changed: 3424 additions & 130 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,3 +3,4 @@ node_modules/
33
__pycache__/
44
.venv/
55
*.py[cod]
6+
ucp/

generate_models.sh

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ rm -rf "$RAW_SCHEMA_DIR"
6363
cp -R "$SCHEMA_DIR" "$RAW_SCHEMA_DIR"
6464

6565
echo "Preprocessing schemas..."
66-
uv run python preprocess_schemas.py
66+
uv run --no-sync python preprocess_schemas.py
6767

6868
echo "Generating Pydantic models from preprocessed schemas..."
6969

@@ -77,9 +77,8 @@ mkdir -p "$OUTPUT_DIR"
7777
# We use --field-constraints to include validation constraints (regex, min/max, etc.)
7878
# We use --reuse-model to collapse structurally identical generated types.
7979
# Note: Formatting is done as a post-processing step.
80-
uv run \
81-
--link-mode=copy \
82-
--extra-index-url https://pypi.org/simple python \
80+
uv run --no-sync \
81+
--link-mode=copy python \
8382
-m datamodel_code_generator \
8483
--input "$SCHEMA_DIR" \
8584
--input-file-type jsonschema \
@@ -99,11 +98,11 @@ uv run \
9998

10099

101100
echo "Post-processing generated models (constraints the generator ignores)..."
102-
uv run python postprocess_models.py || exit 1
101+
uv run --no-sync python postprocess_models.py || exit 1
103102

104103
echo "Formatting generated models..."
105-
uv run ruff format
106-
uv run ruff check --fix "$OUTPUT_DIR"
104+
uv run --no-sync ruff format
105+
uv run --no-sync ruff check --fix "$OUTPUT_DIR"
107106

108107

109108
echo "Done. Models generated in $OUTPUT_DIR"

postprocess_models.py

Lines changed: 49 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -667,7 +667,10 @@ def inject_array_contains(source, alias_name, groups, item_condition=None):
667667
break
668668
if close is None:
669669
return source
670-
out = source[:close] + f", AfterValidator({func_name})" + source[close:]
670+
before = source[:close].rstrip()
671+
if before.endswith(","):
672+
before = before[:-1].rstrip()
673+
out = before + f", AfterValidator({func_name})" + source[close:]
671674
func_src = _build_contains_function(func_name, groups, item_condition)
672675
insert_at = assign_re.search(out).start()
673676
out = out[:insert_at] + func_src + "\n\n" + out[insert_at:]
@@ -1434,6 +1437,48 @@ def _patch_extra_forbid():
14341437
return patched, 0
14351438

14361439

1440+
def _patch_cart_checkout_create_request():
1441+
"""Ensure Checkout in cart_create_request allows cart_id or line_items."""
1442+
path = OUTPUT_DIR / "shopping" / "cart_create_request.py"
1443+
if not path.exists():
1444+
return 0, 0
1445+
source = path.read_text(encoding="utf-8")
1446+
if "_enforce_cart_conversion" in source:
1447+
return 0, 0
1448+
1449+
class_match = re.search(
1450+
r"^class Checkout\(CheckoutCreateRequest\):", source, re.M
1451+
)
1452+
if not class_match:
1453+
return 0, 0
1454+
1455+
validator_code = ''' line_items: list[line_item_create_request.LineItemCreateRequest] | None = None
1456+
1457+
@model_validator(mode="after")
1458+
def _enforce_cart_conversion(self):
1459+
"""Require either cart_id or line_items for checkout creation."""
1460+
if not getattr(self, "cart_id", None) and not getattr(self, "line_items", None):
1461+
raise ValueError("Either cart_id or line_items must be provided")
1462+
return self
1463+
'''
1464+
config_match = re.search(
1465+
r"model_config = ConfigDict\(\s*extra=\"allow\",?\s*\)",
1466+
source[class_match.start() :],
1467+
)
1468+
if config_match:
1469+
insert_pos = class_match.start() + config_match.end()
1470+
source = (
1471+
source[:insert_pos] + "\n" + validator_code + source[insert_pos:]
1472+
)
1473+
source = _ensure_pydantic_import(source, "model_validator")
1474+
path.write_text(source, encoding="utf-8")
1475+
sys.stdout.write(
1476+
f" cart conversion validator on 'Checkout' -> {path}\n"
1477+
)
1478+
return 1, 0
1479+
return 0, 0
1480+
1481+
14371482
def main():
14381483
"""Main entry point to scan schemas and patch generated models."""
14391484
patched_mp, rc_mp = _patch_min_properties()
@@ -1443,6 +1488,7 @@ def main():
14431488
patched_cb, rc_cb = _patch_conditional_bounds()
14441489
patched_ui, rc_ui = _patch_unique_items()
14451490
patched_ef, rc_ef = _patch_extra_forbid()
1491+
patched_cc, rc_cc = _patch_cart_checkout_create_request()
14461492
total = (
14471493
patched_mp
14481494
+ patched_pn
@@ -1451,9 +1497,10 @@ def main():
14511497
+ patched_cb
14521498
+ patched_ui
14531499
+ patched_ef
1500+
+ patched_cc
14541501
)
14551502
sys.stdout.write(f"postprocess: {total} module(s) patched\n")
1456-
return rc_mp or rc_pn or rc_ac or rc_cr or rc_cb or rc_ui or rc_ef
1503+
return rc_mp or rc_pn or rc_ac or rc_cr or rc_cb or rc_ui or rc_ef or rc_cc
14571504

14581505

14591506
if __name__ == "__main__":

preprocess_schemas.py

Lines changed: 53 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -377,21 +377,26 @@ class name like 'Checkout'); fall back to dot-replaced-with-underscore
377377

378378
def get_required_ops(schema):
379379
"""
380-
Scans a schema for the custom 'ucp_request' metadata.
380+
Scans a schema's properties and $defs for custom 'ucp_request' metadata.
381381
Returns a set of operation keys (e.g. {'create', 'update'}) that need distinct models.
382382
"""
383383
ops = set()
384-
properties = schema.get("properties", {})
385-
if not isinstance(properties, dict):
386-
return ops
387-
388-
for data in properties.values():
389-
if isinstance(data, dict):
390-
marker = data.get("ucp_request")
391-
if isinstance(marker, str):
392-
ops.update(["create", "update"]) # Standard shortcut
393-
elif isinstance(marker, dict):
394-
ops.update(marker.keys())
384+
containers = []
385+
if isinstance(schema.get("properties"), dict):
386+
containers.append(schema["properties"])
387+
if isinstance(schema.get("$defs"), dict):
388+
containers.append(schema["$defs"])
389+
390+
for container in containers:
391+
for node in iter_nodes(container):
392+
if not isinstance(node, dict):
393+
continue
394+
marker = node.get("ucp_request")
395+
if marker is not None:
396+
if isinstance(marker, str):
397+
ops.update(["create", "update"]) # Standard shortcut
398+
elif isinstance(marker, dict):
399+
ops.update(marker.keys())
395400
return ops
396401

397402

@@ -527,6 +532,22 @@ def _create_single_variant(
527532
variant, op, file_path, global_variant_requirements
528533
)
529534

535+
# Apply request rules to top-level definitions in $defs
536+
defs = variant.get("$defs", {})
537+
if isinstance(defs, dict):
538+
for node in defs.values():
539+
if isinstance(node, dict) and (
540+
"properties" in node or node.get("type") == "object"
541+
):
542+
_apply_request_rules_to_object(
543+
node, op, file_path, global_variant_requirements
544+
)
545+
546+
# Rewrite all external references across the entire variant tree
547+
rewrite_refs_to_variants(
548+
variant, op, file_path, global_variant_requirements
549+
)
550+
530551
return variant
531552

532553

@@ -595,20 +616,23 @@ def normalize_metadata_schemas(schemas, target_dir):
595616

596617

597618
def extract_external_refs(schema, path):
598-
"""Finds all relative external file references in the schema properties."""
619+
"""Finds all relative external file references in properties and $defs."""
599620
refs = []
600-
props = schema.get("properties", {})
601-
if not isinstance(props, dict):
602-
return refs
603-
604-
for name, data in props.items():
605-
for node in iter_nodes(data):
606-
if isinstance(node, dict) and "$ref" in node:
607-
ref = node["$ref"]
608-
ref_file, _, _ = ref.partition("#")
609-
if ref_file:
610-
abs_path = str((path.parent / ref_file).resolve())
611-
refs.append((name, abs_path))
621+
containers = []
622+
if isinstance(schema.get("properties"), dict):
623+
containers.append(schema["properties"])
624+
if isinstance(schema.get("$defs"), dict):
625+
containers.append(schema["$defs"])
626+
627+
for container in containers:
628+
for name, data in container.items():
629+
for node in iter_nodes(data):
630+
if isinstance(node, dict) and "$ref" in node:
631+
ref = node["$ref"]
632+
ref_file, _, _ = ref.partition("#")
633+
if ref_file:
634+
abs_path = str((path.parent / ref_file).resolve())
635+
refs.append((name, abs_path))
612636
return refs
613637

614638

@@ -629,10 +653,10 @@ def propagate_needs_transitive(variant_needs, schema_refs, schemas):
629653
if child_path not in schemas:
630654
continue
631655

632-
# Only propagate if the property isn't 'omit'ted for this op
633-
data = (
634-
schemas[path].get("properties", {}).get(prop_name, {})
635-
)
656+
# Check properties first, then $defs for any operation override
657+
data = schemas[path].get("properties", {}).get(
658+
prop_name
659+
) or schemas[path].get("$defs", {}).get(prop_name, {})
636660
include, _ = eval_prop_inclusion(
637661
prop_name, data, op, schemas[path].get("required", [])
638662
)

src/ucp_sdk/models/schemas/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,4 @@
1515
# generated by datamodel-codegen
1616
# pylint: disable=all
1717
# pyformat: disable
18+

0 commit comments

Comments
 (0)