Skip to content

Commit ef9a192

Browse files
committed
Fix mypy arg-type errors in generated discriminated union encoders
The codegen for discriminated union TypedDict encoders generates a ternary chain like: encode_Foo(x) if x['kind'] == 'foo' else encode_Bar(x) mypy cannot narrow TypedDict union types through key-based ternary conditions, so every branch (not just the last fallback) receives the full union type. The multi-variant sub-discriminator path (len(oneof_ts) > 1) already had `# type: ignore[arg-type]` on line 306, but the single-variant path (the common case) was missing it. This caused the pid2 codegen CI in ai-infra to fail when the scribe schema added new shape-type discriminated union variants.
1 parent 6c7a537 commit ef9a192

File tree

2 files changed

+3
-2
lines changed

2 files changed

+3
-2
lines changed

src/replit_river/codegen/client.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,7 @@ def flatten_union(tpe: RiverType) -> list[RiverType]:
335335
# encoder_names.add(encoder_name)
336336
typeddict_encoder.append(
337337
f"{render_literal_type(encoder_name)}(x)"
338+
" # type: ignore[arg-type]"
338339
)
339340
typeddict_encoder.append(
340341
f"""

tests/v1/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,9 @@ def encode_NeedsenumobjectInput(
7171
x: "NeedsenumobjectInput",
7272
) -> Any:
7373
return (
74-
encode_NeedsenumobjectInputOneOf_in_first(x)
74+
encode_NeedsenumobjectInputOneOf_in_first(x) # type: ignore[arg-type]
7575
if x["kind"] == "in_first"
76-
else encode_NeedsenumobjectInputOneOf_in_second(x)
76+
else encode_NeedsenumobjectInputOneOf_in_second(x) # type: ignore[arg-type]
7777
)
7878

7979

0 commit comments

Comments
 (0)