fix: scope uniqueItems validators to declaring classes (#61)

* fix: scope uniqueItems validators to declaring classes

* fix: resolve uniqueItems mapping for untitled nested schemas and regenerate models

* style: format code with pre-commit hooks

---------

Co-authored-by: damaz91 <federico.damato91@gmail.com>

Author: horaceewang-TT

postprocess_models.py

@@ -248,6 +248,12 @@ def _alias_name(title):
     return "".join(title.split())
 
 
+def _to_camel_case(string):
+    """Convert a string (snake, kebab, space-separated) to CamelCase."""
+    parts = re.split(r"[^a-zA-Z0-9]", string)
+    return "".join(p.capitalize() for p in parts if p)
+
+
 def _snake_name(name):
     """CamelCase alias -> snake_case suffix for a unique function name."""
     return re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
@@ -346,65 +352,84 @@ def inject_array_contains(source, alias_name, groups):
     return _ensure_pydantic_import(out, "AfterValidator")
 
 
-def _iter_nodes(root):
-    """Yield every dict/list node in a JSON tree (cycle-safe)."""
-    stack = [root]
-    seen = {id(root)}
-    while stack:
-        cur = stack.pop()
-        yield cur
-        if isinstance(cur, dict):
-            children = cur.values()
-        elif isinstance(cur, list):
-            children = cur
-        else:
-            children = ()
-        for child in children:
-            if isinstance(child, (dict, list)) and id(child) not in seen:
-                seen.add(id(child))
-                stack.append(child)
-
-
 def find_unique_items_fields(schema_dir):
-    """Collect property names whose array value carries ``uniqueItems``.
+    """Map generated class names to fields carrying ``uniqueItems``.
 
-    Walks every schema (root and nested) for object properties declared as an
-    array with ``uniqueItems: true``. Returns the set of property names so the
-    injector can locate the matching generated list fields by name.
+    A schema node needs a title so its constraint can be associated with a
+    generated class. Untitled nodes are resolved using their property path.
     """
-    fields = set()
+    fields_by_class = {}
+
+    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"])
+
+        props = node.get("properties")
+        if isinstance(props, dict):
+            for name, prop in props.items():
+                if not isinstance(prop, dict):
+                    continue
+
+                if prop.get("uniqueItems") is True and (
+                    prop.get("type") == "array" or "items" in prop
+                ):
+                    if current_class_name is None:
+                        sys.stderr.write(
+                            f"  ! {path_str}: uniqueItems field '{name}' "
+                            "belongs to an untitled object; cannot map to a class\n"
+                        )
+                        continue
+                    fields_by_class.setdefault(current_class_name, set()).add(
+                        name
+                    )
+
+                # Recurse into properties
+                next_class_name = (
+                    _to_camel_case(name) if current_class_name else None
+                )
+                walk(prop, next_class_name, path_str)
+
+        # Recurse into $defs
+        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)
+
+        # Recurse into combinators (allOf, anyOf, oneOf)
+        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
-        for node in _iter_nodes(schema):
-            if not isinstance(node, dict):
-                continue
-            props = node.get("properties")
-            if not isinstance(props, dict):
-                continue
-            for name, prop in props.items():
-                if (
-                    isinstance(prop, dict)
-                    and prop.get("uniqueItems") is True
-                    and (prop.get("type") == "array" or "items" in prop)
-                ):
-                    fields.add(name)
-    return fields
 
+        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 fields_by_class
 
-def inject_unique_items(source, unique_fields):
+
+def inject_unique_items(source, unique_fields_by_class):
     """Inject uniqueness validators for list fields declared ``uniqueItems``.
 
-    Scans each generated class for list-typed fields whose name is in
-    ``unique_fields`` and appends a ``field_validator`` to the class body.
+    A validator is added only when both the generated class name and list
+    field name match the scoped schema constraints.
     """
-    if not unique_fields:
+    if not unique_fields_by_class:
         return source
-    class_re = re.compile(r"^class \w+\(", re.M)
+    class_re = re.compile(r"^class (\w+)\(", re.M)
     matches = list(class_re.finditer(source))
     if not matches:
         return source
@@ -413,6 +438,9 @@ def inject_unique_items(source, unique_fields):
     # Process from the last class to the first so earlier insert offsets
     # (computed against the original source) stay valid as text is appended.
     for match in reversed(matches):
+        unique_fields = unique_fields_by_class.get(match.group(1), set())
+        if not unique_fields:
+            continue
         body_start = match.end()
         tail = re.compile(r"^\S", re.M)
         end_match = tail.search(source, body_start)
@@ -549,21 +577,26 @@ def _patch_array_contains():
 
 def _patch_unique_items():
     """Inject uniqueItems validators; return (patched_count, exit_code)."""
-    unique_fields = find_unique_items_fields(SCHEMA_DIR)
-    if not unique_fields:
+    unique_fields_by_class = find_unique_items_fields(SCHEMA_DIR)
+    if not unique_fields_by_class:
         sys.stdout.write("postprocess: no uniqueItems constraints found\n")
         return 0, 0
     unique_patched = 0
     touched = []
     for path in sorted(OUTPUT_DIR.rglob("*.py")):
         source = path.read_text(encoding="utf-8")
-        updated = inject_unique_items(source, unique_fields)
+        updated = inject_unique_items(source, unique_fields_by_class)
         if updated != source:
             path.write_text(updated, encoding="utf-8")
             unique_patched += 1
             touched.append(path)
+    labels = sorted(
+        f"{class_name}.{field}"
+        for class_name, fields in unique_fields_by_class.items()
+        for field in fields
+    )
     sys.stdout.write(
-        f"  uniqueItems fields {sorted(unique_fields)} -> "
+        f"  uniqueItems fields {labels} -> "
         f"{unique_patched} module(s) patched"
         f" ({', '.join(str(t) for t in touched) or 'none'})\n"
     )

tests/test_codegen_pipeline.py

@@ -912,6 +912,7 @@ class UniqueItemsInjectorTest(unittest.TestCase):
     """The uniqueItems post-generation injector's own behavior."""
 
     SCHEMA_TREE = {
+        "title": "First",
         "properties": {
             "tags": {
                 "type": "array",
@@ -930,7 +931,7 @@ class UniqueItemsInjectorTest(unittest.TestCase):
                     }
                 },
             },
-        }
+        },
     }
 
     MODULE = (
@@ -955,6 +956,7 @@ class UniqueItemsInjectorTest(unittest.TestCase):
         "    model_config = ConfigDict(\n"
         '        extra="allow",\n'
         "    )\n"
+        "    tags: list[str] | None = None\n"
         "    count: list[int] | None = None\n"
     )
 
@@ -963,7 +965,7 @@ def test_find_unique_items_fields_walks_nested_properties(self) -> None:
         with tempfile.TemporaryDirectory() as tmp:
             (Path(tmp) / "schema.json").write_text(json.dumps(self.SCHEMA_TREE))
             fields = postprocess_models.find_unique_items_fields(Path(tmp))
-        self.assertEqual(fields, {"tags", "codes"})
+        self.assertEqual(fields, {"First": {"tags"}, "Nested": {"codes"}})
 
     def test_find_unique_items_fields_ignores_false_and_non_arrays(
         self,
@@ -982,33 +984,43 @@ def test_find_unique_items_fields_ignores_false_and_non_arrays(
         with tempfile.TemporaryDirectory() as tmp:
             (Path(tmp) / "s.json").write_text(json.dumps(schema))
             fields = postprocess_models.find_unique_items_fields(Path(tmp))
-        self.assertEqual(fields, set())
+        self.assertEqual(fields, {})
 
     def test_inject_targets_matching_list_fields_only(self) -> None:
-        """Only list fields named in the set get a validator."""
-        out = postprocess_models.inject_unique_items(self.MODULE, {"tags"})
+        """Only the declaring class's matching list field gets a validator."""
+        out = postprocess_models.inject_unique_items(
+            self.MODULE, {"First": {"tags"}}
+        )
         self.assertIn("field_validator", out)
         self.assertIn("_enforce_unique_items_tags", out)
+        self.assertEqual(out.count("def _enforce_unique_items_tags("), 1)
         self.assertNotIn("_enforce_unique_items_name", out)
         self.assertNotIn("_enforce_unique_items_count", out)
 
     def test_inject_no_match_leaves_source_unchanged(self) -> None:
         """No matching list field means the module is untouched."""
         self.assertEqual(
-            postprocess_models.inject_unique_items(self.MODULE, {"missing"}),
+            postprocess_models.inject_unique_items(
+                self.MODULE, {"First": {"missing"}}
+            ),
             self.MODULE,
         )
 
     def test_injection_is_idempotent(self) -> None:
         """Re-running the injector changes nothing."""
-        once = postprocess_models.inject_unique_items(self.MODULE, {"tags"})
-        twice = postprocess_models.inject_unique_items(once, {"tags"})
+        unique_fields = {"First": {"tags"}}
+        once = postprocess_models.inject_unique_items(
+            self.MODULE, unique_fields
+        )
+        twice = postprocess_models.inject_unique_items(once, unique_fields)
         self.assertEqual(once, twice)
 
     @unittest.skipUnless(HAVE_SDK, "executing the module needs pydantic")
     def test_injected_validator_rejects_duplicates(self) -> None:
         """The injected field_validator enforces uniqueness at runtime."""
-        out = postprocess_models.inject_unique_items(self.MODULE, {"tags"})
+        out = postprocess_models.inject_unique_items(
+            self.MODULE, {"First": {"tags"}}
+        )
         namespace: dict = {}
         exec(compile(out, "<injected>", "exec"), namespace)  # noqa: S102
         first = namespace["First"]