fix: enforce custom totals display text (#75)

Author: FanouZeng-TT

postprocess_models.py

@@ -488,17 +488,66 @@ def find_array_contains_constraints(schema_dir):
         groups = _extract_contains_groups(schema, path)
         if not groups:
             continue
+        item_condition = _extract_item_required_condition(schema)
         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}
+        found[path.stem] = {
+            "title": title,
+            "groups": groups,
+            "item_condition": item_condition,
+        }
     return found
 
 
+def _extract_item_required_condition(schema):
+    """Read a simple array-item ``not.enum`` + ``then.required`` rule."""
+    items = schema.get("items")
+    if not isinstance(items, dict):
+        return None
+    nodes = [items]
+    nodes.extend(
+        node for node in items.get("allOf", []) if isinstance(node, dict)
+    )
+    for node in nodes:
+        condition = node.get("if")
+        consequence = node.get("then")
+        if not isinstance(condition, dict) or not isinstance(consequence, dict):
+            continue
+        props = condition.get("properties")
+        required = consequence.get("required")
+        if not isinstance(props, dict) or len(props) != 1:
+            continue
+        field, predicate = next(iter(props.items()))
+        if (
+            not isinstance(predicate, dict)
+            or set(node) != {"if", "then"}
+            or set(condition) != {"properties", "required"}
+            or set(consequence) != {"required"}
+        ):
+            continue
+        excluded = predicate.get("not")
+        values = excluded.get("enum") if isinstance(excluded, dict) else None
+        if (
+            condition.get("required") == [field]
+            and set(predicate) == {"not"}
+            and isinstance(excluded, dict)
+            and set(excluded) == {"enum"}
+            and isinstance(values, list)
+            and values
+            and all(isinstance(value, str) for value in values)
+            and isinstance(required, list)
+            and required
+            and all(isinstance(name, str) for name in required)
+        ):
+            return {"field": field, "excluded": values, "required": required}
+    return None
+
+
 def _alias_name(title):
     """Derive the generated alias name from a schema title (drop spaces)."""
     return "".join(title.split())
@@ -530,7 +579,7 @@ def _predicate_expr(pairs):
     return " and ".join(parts)
 
 
-def _build_contains_function(func_name, groups):
+def _build_contains_function(func_name, groups, item_condition=None):
     """Render the module-level ``AfterValidator`` counting function."""
     lines = [
         f"def {func_name}(value):",
@@ -565,11 +614,28 @@ def _build_contains_function(func_name, groups):
                 f'            "matching {desc} (schema maxContains={maximum})"',
                 "        )",
             ]
+    if item_condition:
+        field = item_condition["field"]
+        lines += [
+            f"    _excluded = {item_condition['excluded']!r}",
+            "    for _item in value:",
+            f"        _actual = (_item.get({field!r}) if isinstance(_item, dict) ",
+            f"                   else getattr(_item, {field!r}, None))",
+            "        if _actual in _excluded:",
+            "            continue",
+        ]
+        for required in item_condition["required"]:
+            lines += [
+                f"        if isinstance(_item, dict) and {required!r} not in _item:",
+                f'            raise ValueError("Field {required!r} is required for custom {field}")',
+                f"        if not isinstance(_item, dict) and {required!r} not in _item.model_fields_set:",
+                f'            raise ValueError("Field {required!r} is required for custom {field}")',
+            ]
     lines.append("    return value")
     return "\n".join(lines) + "\n"
 
 
-def inject_array_contains(source, alias_name, groups):
+def inject_array_contains(source, alias_name, groups, item_condition=None):
     """Thread an ``AfterValidator`` into ``alias_name``'s alias metadata.
 
     Array roots are emitted as ``NAME = TypeAliasType("NAME", Annotated[...])``,
@@ -602,7 +668,7 @@ def inject_array_contains(source, alias_name, groups):
     if close is None:
         return source
     out = source[:close] + f", AfterValidator({func_name})" + source[close:]
-    func_src = _build_contains_function(func_name, groups)
+    func_src = _build_contains_function(func_name, groups, item_condition)
     insert_at = assign_re.search(out).start()
     out = out[:insert_at] + func_src + "\n\n" + out[insert_at:]
     return _ensure_pydantic_import(out, "AfterValidator")
@@ -1082,11 +1148,20 @@ def _array_contains_targets():
             None,
         )
         if origin is not None:
-            targets[info["title"]] = raw[origin]["groups"]
+            targets[info["title"]] = {
+                "groups": raw[origin]["groups"],
+                "item_condition": raw[origin]["item_condition"],
+            }
     # 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"])
+        targets.setdefault(
+            info["title"],
+            {
+                "groups": info["groups"],
+                "item_condition": info["item_condition"],
+            },
+        )
     return targets
 
 
@@ -1133,7 +1208,8 @@ def _patch_array_contains():
         sys.stdout.write("postprocess: no array contains constraints found\n")
         return 0, 0
     patched = 0
-    for title, groups in sorted(targets.items()):
+    for title, constraint in sorted(targets.items()):
+        groups = constraint["groups"]
         alias = _alias_name(title)
         hits = []
         for path in sorted(OUTPUT_DIR.rglob("*.py")):
@@ -1142,7 +1218,9 @@ def _patch_array_contains():
                 rf"^{re.escape(alias)} = TypeAliasType\(", source, re.M
             ):
                 continue
-            updated = inject_array_contains(source, alias, groups)
+            updated = inject_array_contains(
+                source, alias, groups, constraint["item_condition"]
+            )
             if updated != source:
                 path.write_text(updated, encoding="utf-8")
                 patched += 1

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

@@ -140,6 +140,30 @@ def _enforce_contains_totals(value):
             "Array must contain at most 1 entry "
             "matching type=='total' (schema maxContains=1)"
         )
+    _excluded = [
+        "subtotal",
+        "items_discount",
+        "discount",
+        "fulfillment",
+        "tax",
+        "fee",
+        "total",
+    ]
+    for _item in value:
+        _actual = (
+            _item.get("type")
+            if isinstance(_item, dict)
+            else getattr(_item, "type", None)
+        )
+        if _actual in _excluded:
+            continue
+        if isinstance(_item, dict) and "display_text" not in _item:
+            raise ValueError("Field 'display_text' is required for custom type")
+        if (
+            not isinstance(_item, dict)
+            and "display_text" not in _item.model_fields_set
+        ):
+            raise ValueError("Field 'display_text' is required for custom type")
     return value
 
 

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

@@ -68,6 +68,30 @@ def _enforce_contains_totals_create_request(value):
             "Array must contain at most 1 entry "
             "matching type=='total' (schema maxContains=1)"
         )
+    _excluded = [
+        "subtotal",
+        "items_discount",
+        "discount",
+        "fulfillment",
+        "tax",
+        "fee",
+        "total",
+    ]
+    for _item in value:
+        _actual = (
+            _item.get("type")
+            if isinstance(_item, dict)
+            else getattr(_item, "type", None)
+        )
+        if _actual in _excluded:
+            continue
+        if isinstance(_item, dict) and "display_text" not in _item:
+            raise ValueError("Field 'display_text' is required for custom type")
+        if (
+            not isinstance(_item, dict)
+            and "display_text" not in _item.model_fields_set
+        ):
+            raise ValueError("Field 'display_text' is required for custom type")
     return value
 
 

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

@@ -68,6 +68,30 @@ def _enforce_contains_totals_update_request(value):
             "Array must contain at most 1 entry "
             "matching type=='total' (schema maxContains=1)"
         )
+    _excluded = [
+        "subtotal",
+        "items_discount",
+        "discount",
+        "fulfillment",
+        "tax",
+        "fee",
+        "total",
+    ]
+    for _item in value:
+        _actual = (
+            _item.get("type")
+            if isinstance(_item, dict)
+            else getattr(_item, "type", None)
+        )
+        if _actual in _excluded:
+            continue
+        if isinstance(_item, dict) and "display_text" not in _item:
+            raise ValueError("Field 'display_text' is required for custom type")
+        if (
+            not isinstance(_item, dict)
+            and "display_text" not in _item.model_fields_set
+        ):
+            raise ValueError("Field 'display_text' is required for custom type")
     return value
 
 

tests/test_codegen_pipeline.py

@@ -1293,6 +1293,27 @@ def test_create_request_variant_enforces_both_bounds(self):
     def test_update_request_variant_enforces_both_bounds(self):
         self._assert_matrix(TotalsUpdateRequest)
 
+    def test_custom_type_requires_display_text(self):
+        base = [self.SUBTOTAL, self.TOTAL]
+        for alias in (Totals, TotalsCreateRequest, TotalsUpdateRequest):
+            adapter = TypeAdapter(alias)
+            with self.subTest(model=alias.__name__):
+                with self.assertRaisesRegex(ValidationError, "display_text"):
+                    adapter.validate_python(
+                        base + [{"type": "surcharge", "amount": 5}]
+                    )
+                adapter.validate_python(base + [{"type": "tax", "amount": 5}])
+                adapter.validate_python(
+                    base
+                    + [
+                        {
+                            "type": "surcharge",
+                            "amount": 5,
+                            "display_text": "Surcharge",
+                        }
+                    ]
+                )
+
     def test_missing_total_names_the_total_rule(self):
         # A subtotal-only array must fail specifically on the total rule.
         with self.assertRaisesRegex(ValidationError, "total"):
@@ -1328,11 +1349,30 @@ class ArrayContainsInjectorTest(unittest.TestCase):
         {"pairs": [("type", "total")], "min": 1, "max": 1},
     ]
 
+    ITEM_CONDITION = {
+        "field": "type",
+        "excluded": ["subtotal", "total"],
+        "required": ["display_text"],
+    }
+
     def test_scan_reads_both_contains_from_allof_branches(self):
         # The pristine totals.json shape: two allOf contains branches.
         schema = {
             "title": "Totals",
             "type": "array",
+            "items": {
+                "allOf": [
+                    {
+                        "if": {
+                            "properties": {
+                                "type": {"not": {"enum": ["subtotal", "total"]}}
+                            },
+                            "required": ["type"],
+                        },
+                        "then": {"required": ["display_text"]},
+                    }
+                ]
+            },
             "allOf": [
                 {
                     "contains": {"properties": {"type": {"const": "subtotal"}}},
@@ -1357,6 +1397,10 @@ def test_scan_reads_both_contains_from_allof_branches(self):
             [g["pairs"] for g in found["totals"]["groups"]],
             [[("type", "subtotal")], [("type", "total")]],
         )
+        self.assertEqual(
+            found["totals"]["item_condition"],
+            self.ITEM_CONDITION,
+        )
 
     def test_scan_reads_root_level_single_contains(self):
         # A root-level (non-allOf) contains still yields one group.
@@ -1435,6 +1479,31 @@ def test_injected_validator_enforces_both_bounds(self):
                 adapter.validate_python(bad)
         adapter.validate_python([sub, tot])
 
+    @unittest.skipUnless(HAVE_SDK, "executing the module needs pydantic")
+    def test_injected_validator_requires_custom_display_text(self):
+        out = postprocess_models.inject_array_contains(
+            self.MODULE, "Totals", self.GROUPS, self.ITEM_CONDITION
+        )
+        namespace: dict = {}
+        exec(compile(out, "<injected>", "exec"), namespace)  # noqa: S102
+        adapter = TypeAdapter(namespace["Totals"])
+        base = [
+            {"type": "subtotal", "amount": 1},
+            {"type": "total", "amount": 1},
+        ]
+        with self.assertRaisesRegex(ValidationError, "display_text"):
+            adapter.validate_python(base + [{"type": "surcharge", "amount": 1}])
+        adapter.validate_python(
+            base
+            + [
+                {
+                    "type": "surcharge",
+                    "amount": 1,
+                    "display_text": "Surcharge",
+                }
+            ]
+        )
+
 
 class UniqueItemsInjectorTest(unittest.TestCase):
     """The uniqueItems post-generation injector's own behavior."""