fix(types): derive generated account union arms

Author: bokelley

scripts/post_generate_fixes.py

@@ -1545,14 +1545,58 @@ def expose_account_reference_union_fields() -> None:
     request, nested-input, response, and canonical-clone paths all expose the
     same concrete arm types without import-time Pydantic patching.
     """
-    pattern = re.compile(r"\b(account_ref(?:_\d+)?)\.AccountReference\b(?![12])")
+    account_ref_source = OUTPUT_DIR / "core" / "account_ref.py"
+    if not account_ref_source.exists():
+        print("  account reference model not found (skipping union-field fix)")
+        return
+
+    tree = ast.parse(account_ref_source.read_text())
+    wrapper = next(
+        (
+            node
+            for node in tree.body
+            if isinstance(node, ast.ClassDef) and node.name == "AccountReference"
+        ),
+        None,
+    )
+    if wrapper is None:
+        raise RuntimeError("generated account_ref.py has no AccountReference wrapper")
+
+    root_base = next(
+        (
+            base
+            for base in wrapper.bases
+            if isinstance(base, ast.Subscript)
+            and isinstance(base.value, ast.Name)
+            and base.value.id == "RootModel"
+        ),
+        None,
+    )
+    if root_base is None:
+        raise RuntimeError("generated AccountReference has no RootModel union base")
+
+    def union_arm_names(node: ast.expr) -> list[str]:
+        if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
+            return [*union_arm_names(node.left), *union_arm_names(node.right)]
+        if isinstance(node, ast.Name):
+            return [node.id]
+        raise RuntimeError(
+            "generated AccountReference has an unsupported union expression: "
+            f"{ast.unparse(node)}"
+        )
+
+    arm_names = union_arm_names(root_base.slice)
+    if len(arm_names) < 2 or len(set(arm_names)) != len(arm_names):
+        raise RuntimeError(f"generated AccountReference has invalid union arms: {arm_names!r}")
+
+    pattern = re.compile(r"\b(account_ref(?:_\d+)?)\.AccountReference\b(?!\d)")
     total_files = 0
     total_fields = 0
 
     for py_file in sorted(OUTPUT_DIR.rglob("*.py")):
         source = py_file.read_text()
         fixed, replacements = pattern.subn(
-            r"\1.AccountReference1 | \1.AccountReference2",
+            lambda match: " | ".join(f"{match.group(1)}.{arm_name}" for arm_name in arm_names),
             source,
         )
         if not replacements:

tests/test_code_generation.py

@@ -378,8 +378,15 @@ def test_post_generate_exposes_account_reference_union_fields_idempotently(tmp_p
     from scripts import post_generate_fixes
 
     generated_dir = tmp_path / "generated_poc"
+    account_ref = generated_dir / "core" / "account_ref.py"
+    account_ref.parent.mkdir(parents=True)
+    account_ref.write_text(
+        "from pydantic import RootModel\n\n"
+        "class AccountReference(RootModel["
+        "AccountReference1 | AccountReference2 | AccountReference3]):\n"
+        "    pass\n"
+    )
     target = generated_dir / "sample_request.py"
-    target.parent.mkdir(parents=True)
     target.write_text(
         "account: account_ref.AccountReference | None\n"
         "accounts: list[account_ref_1.AccountReference]\n"
@@ -390,8 +397,10 @@ def test_post_generate_exposes_account_reference_union_fields_idempotently(tmp_p
     post_generate_fixes.expose_account_reference_union_fields()
 
     assert target.read_text() == (
-        "account: account_ref.AccountReference1 | account_ref.AccountReference2 | None\n"
-        "accounts: list[account_ref_1.AccountReference1 | account_ref_1.AccountReference2]\n"
+        "account: account_ref.AccountReference1 | account_ref.AccountReference2 | "
+        "account_ref.AccountReference3 | None\n"
+        "accounts: list[account_ref_1.AccountReference1 | account_ref_1.AccountReference2 | "
+        "account_ref_1.AccountReference3]\n"
     )