codegen: emit guards for every anyOf variant to fix mypy union-attr on array-containing unions (#182)

Why
===

The encoder generator for non-discriminated anyOf unions emits a chain
of ternary expressions, with the last variant historically rendered as
the unguarded `else` branch. That works for simple unions like `object |
str | list` (mypy can negative-narrow `x` to `list` in the final
branch), but it breaks for deeper unions where the array variant is
last, e.g.

```
str | float | bool | list[scalar] | None
```

When mypy fails to fully narrow `x` to `list[...]` through the prior
`isinstance` checks (`isinstance(x, (int, float))` plus `bool`
subclassing `int` make this tricky), it complains that scalar items of
the union have no `__iter__` attribute:

```
error: Item "float" of "str | float | bool | None | list[...]"
    has no attribute "__iter__" (not iterable)  [union-attr]
error: Item "bool"  of ... has no attribute "__iter__"  [union-attr]
error: Item "object" of ... has no attribute "__iter__"  [union-attr]
```

This is the exact failure that has been blocking ai-infra's
`codegen-latest-pid2-schema.yml` auto-update workflow since 2026-05-04,
when repl-it-web#78355 widened
`agentToolPostgreSQL.executeSqlCommand.params` from a flat scalar union
to `array<scalar | array<scalar>>`. Every run since has failed on the
regenerated `executeSqlCommand.py` at the `for y in x` iteration inside
`encode_ExecutesqlcommandInputParams`. The committed pid2 client in
ai-infra has been kept current by hand (see replit/ai-infra#12813), but
the bot has been red for ~2.5 weeks.

What changed
============

`src/replit_river/codegen/client.py`: in the non-discriminated-anyOf
branch of `encode_type`, emit an explicit `isinstance` / `is None` guard
for every entry in `encoder_parts` — including the last one — and append
a `cast(Any, x)` fallback. mypy no longer has to negative-narrow into
the iterating branch, so deep unions with an array variant lint cleanly.
`Any` and `cast` are already part of `FILE_HEADER` so no import
bookkeeping changes.

Concretely, for the failing executeSqlCommand schema, the encoder now
ends with:

```python
return (
    x if isinstance(x, str)
    else x if isinstance(x, (int, float))
    else x if isinstance(x, bool)
    else None if x is None
    else [encode_..._AnyOf_4(y) for y in x]
        if isinstance(x, list)
        else cast(Any, x)
)
```

Test plan
=========

- Existing `tests/v1/codegen/snapshot/test_anyof_mixed.py` snapshot
updated to show the new `if isinstance(x, list) else cast(Any, x)` tail
on its `obj | str | list[str]` encoder (the change is additive — the
runtime behavior is unchanged).
- New snapshot test
`tests/v1/codegen/snapshot/test_anyof_array_in_union.py` added with a
schema that mirrors `executeSqlCommand.params` (`array<scalar |
array<scalar>>`) and locks in the fixed output. This is the regression
test for ai-infra's CI failure.
- `uv run pytest` is green (67 passed, including all v1 and v2 codegen
tests).
- `make lint` is clean apart from a pre-existing `pyright` `grpc` import
error in `tests/v1/test_communication.py` that also fails on `main`
(unrelated).
- End-to-end verification against ai-infra: pointed ai-infra's
`./pkgs/pid2_client/scripts/generate.sh` at this branch via
`RIVER_CODEGEN_PATH=/tmp/opencode/river-python` and reran the full lint
pipeline that the auto-update workflow runs in CI; `[mypy] completed in
15.19s` and the script exited `OK.` instead of the historical
`union-attr` failure.

Once this is released (e.g. `v0.17.20`) ai-infra can bump `replit-river`
in `pkgs/pid2_client/pyproject.toml` and the auto-update workflow will
start producing green PRs again.

~ written by Zerg 👾
([ascendant-goliath-6d2f](https://zerg.zergrush.dev/chat?id=ascendant-goliath-6d2f))

Author: kush-replit

src/replit_river/codegen/client.py

@@ -457,16 +457,35 @@ def {_field_name}(
                         case _:
                             encoder_parts.append((None, "x"))
 
-        # Build the ternary chain from encoder_parts
+        # Build the ternary chain from encoder_parts.
+        #
+        # Every entry that has a `type_check` (isinstance / `x is None`) gets
+        # its own guard, including the last one. Falling off the end means the
+        # input did not match any declared anyOf variant, which should not
+        # happen for a well-formed value; we emit a `cast(Any, x)` so mypy
+        # doesn't try to narrow the value through the chain.
+        #
+        # Previously the last entry was emitted unconditionally as the `else`
+        # branch. That works for simple unions (object | str | list), but
+        # breaks down when the last variant's encoder requires iteration
+        # (e.g. `[encode_X(y) for y in x]` when the variant is an array) and
+        # mypy fails to fully narrow `x` through the prior `isinstance`
+        # checks. The unguarded final branch then triggers `union-attr`
+        # errors like "Item 'float' has no attribute '__iter__'".
         typeddict_encoder = list[str]()
-        for i, (type_check, encoder_expr) in enumerate(encoder_parts):
-            is_last = i == len(encoder_parts) - 1
-            if is_last or type_check is None:
-                # Last item or no type check - just the expression
+        has_unguarded_terminal = False
+        for type_check, encoder_expr in encoder_parts:
+            if type_check is None:
+                # No type check available — emit the bare expression and stop;
+                # nothing after it could be reached anyway.
                 typeddict_encoder.append(encoder_expr)
-            else:
-                # Add expression with type check
-                typeddict_encoder.append(f"{encoder_expr} if {type_check} else")
+                has_unguarded_terminal = True
+                break
+            typeddict_encoder.append(f"{encoder_expr} if {type_check} else")
+        if not has_unguarded_terminal and encoder_parts:
+            # Unreachable in practice (every declared variant was guarded
+            # above), but mypy needs a concrete final expression.
+            typeddict_encoder.append("cast(Any, x)")
         if permit_unknown_members:
             union = _make_open_union_type_expr(any_of)
         else:

tests/v1/codegen/snapshot/snapshots/test_anyof_array_in_union/__init__.py

@@ -0,0 +1,13 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from pydantic import BaseModel
+from typing import Literal
+
+import replit_river as river
+
+
+from .test_service import Test_ServiceService
+
+
+class AnyOfArrayInUnionClient:
+    def __init__(self, client: river.Client[Literal[None]]):
+        self.test_service = Test_ServiceService(client)

tests/v1/codegen/snapshot/snapshots/test_anyof_array_in_union/test_service/__init__.py

@@ -0,0 +1,42 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from collections.abc import AsyncIterable, AsyncIterator
+from typing import Any
+import datetime
+
+from pydantic import TypeAdapter
+
+from replit_river.error_schema import RiverError, RiverErrorTypeAdapter
+import replit_river as river
+
+
+from .exec_sql_method import (
+    Exec_Sql_MethodInput,
+    Exec_Sql_MethodOutput,
+    Exec_Sql_MethodOutputTypeAdapter,
+    encode_Exec_Sql_MethodInput,
+    encode_Exec_Sql_MethodInputParams,
+)
+
+
+class Test_ServiceService:
+    def __init__(self, client: river.Client[Any]):
+        self.client = client
+
+    async def exec_sql_method(
+        self,
+        input: Exec_Sql_MethodInput,
+        timeout: datetime.timedelta,
+    ) -> Exec_Sql_MethodOutput:
+        return await self.client.send_rpc(
+            "test_service",
+            "exec_sql_method",
+            input,
+            encode_Exec_Sql_MethodInput,
+            lambda x: Exec_Sql_MethodOutputTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            lambda x: RiverErrorTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            timeout,
+        )

tests/v1/codegen/snapshot/snapshots/test_anyof_array_in_union/test_service/exec_sql_method.py

@@ -0,0 +1,82 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from collections.abc import AsyncIterable, AsyncIterator
+import datetime
+from typing import (
+    Any,
+    Literal,
+    Mapping,
+    NotRequired,
+    TypedDict,
+    cast,
+)
+from typing_extensions import Annotated
+
+from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
+from replit_river.error_schema import RiverError
+from replit_river.client import (
+    RiverUnknownError,
+    translate_unknown_error,
+    RiverUnknownValue,
+    translate_unknown_value,
+)
+
+import replit_river as river
+
+
+Exec_Sql_MethodInputParamsAnyOf_4 = str | float | bool | None
+
+
+def encode_Exec_Sql_MethodInputParamsAnyOf_4(
+    x: "Exec_Sql_MethodInputParamsAnyOf_4",
+) -> Any:
+    return x
+
+
+Exec_Sql_MethodInputParams = (
+    str | float | bool | list[Exec_Sql_MethodInputParamsAnyOf_4] | None
+)
+
+
+def encode_Exec_Sql_MethodInputParams(x: "Exec_Sql_MethodInputParams") -> Any:
+    return (
+        x
+        if isinstance(x, str)
+        else x
+        if isinstance(x, (int, float))
+        else x
+        if isinstance(x, bool)
+        else None
+        if x is None
+        else [encode_Exec_Sql_MethodInputParamsAnyOf_4(y) for y in x]
+        if isinstance(x, list)
+        else cast(Any, x)
+    )
+
+
+def encode_Exec_Sql_MethodInput(
+    x: "Exec_Sql_MethodInput",
+) -> Any:
+    return {
+        k: v
+        for (k, v) in (
+            {
+                "params": [encode_Exec_Sql_MethodInputParams(y) for y in x["params"]]
+                if "params" in x and x["params"] is not None
+                else None,
+            }
+        ).items()
+        if v is not None
+    }
+
+
+class Exec_Sql_MethodInput(TypedDict):
+    params: NotRequired[list[Exec_Sql_MethodInputParams] | None]
+
+
+class Exec_Sql_MethodOutput(BaseModel):
+    ok: bool
+
+
+Exec_Sql_MethodOutputTypeAdapter: TypeAdapter[Exec_Sql_MethodOutput] = TypeAdapter(
+    Exec_Sql_MethodOutput
+)

tests/v1/codegen/snapshot/snapshots/test_anyof_mixed_types/test_service/anyof_mixed_method.py

@@ -66,6 +66,8 @@ def encode_Anyof_Mixed_MethodInputRun_Command(
         else x
         if isinstance(x, str)
         else list(x)
+        if isinstance(x, list)
+        else cast(Any, x)
     )
 
 

tests/v1/codegen/snapshot/test_anyof_array_in_union.py

@@ -0,0 +1,31 @@
+from pytest_snapshot.plugin import Snapshot
+
+from tests.fixtures.codegen_snapshot_fixtures import validate_codegen
+
+
+async def test_anyof_array_in_union(snapshot: Snapshot) -> None:
+    """Test codegen for an array field whose item type is a non-discriminated
+    anyOf union that itself contains an `array` variant.
+
+    Concretely this mirrors the PostgreSQL `executeSqlCommand.params` schema:
+    `array<scalar | array<scalar>>`. The inner union encoder ends in an
+    iteration over `x` (for the array variant), and historically that branch
+    was emitted as the unguarded `else` of a ternary chain. When mypy failed
+    to fully narrow `x` to `list[...]` through the preceding `isinstance`
+    checks, it complained that scalar items of the union have no
+    `__iter__` attribute (`union-attr`).
+
+    The fix emits an explicit `isinstance(x, list)` guard for the array
+    branch and a `cast(Any, x)` fallback, so mypy never has to negative-
+    narrow into the iterating branch.
+    """
+    validate_codegen(
+        snapshot=snapshot,
+        snapshot_dir="tests/v1/codegen/snapshot/snapshots",
+        read_schema=lambda: open(
+            "tests/v1/codegen/types/anyof_array_in_union_schema.json"
+        ),
+        target_path="test_anyof_array_in_union",
+        client_name="AnyOfArrayInUnionClient",
+        protocol_version="v1.1",
+    )

tests/v1/codegen/types/anyof_array_in_union_schema.json

@@ -0,0 +1,49 @@
+{
+  "services": {
+    "test_service": {
+      "procedures": {
+        "exec_sql_method": {
+          "input": {
+            "type": "object",
+            "properties": {
+              "params": {
+                "description": "Parameterized query values. Each entry is either a scalar or an array of scalars (for ANY($1::text[]) etc.).",
+                "type": "array",
+                "items": {
+                  "anyOf": [
+                    { "type": "string" },
+                    { "type": "number" },
+                    { "type": "boolean" },
+                    { "type": "null" },
+                    {
+                      "type": "array",
+                      "items": {
+                        "anyOf": [
+                          { "type": "string" },
+                          { "type": "number" },
+                          { "type": "boolean" },
+                          { "type": "null" }
+                        ]
+                      }
+                    }
+                  ]
+                }
+              }
+            }
+          },
+          "output": {
+            "type": "object",
+            "properties": {
+              "ok": { "type": "boolean" }
+            },
+            "required": ["ok"]
+          },
+          "errors": {
+            "not": {}
+          },
+          "type": "rpc"
+        }
+      }
+    }
+  }
+}