From ab95ace7abd3d85c3c8e49404c116af19722842c Mon Sep 17 00:00:00 2001 From: Darsh Patel Date: Thu, 5 Mar 2026 13:45:50 -0800 Subject: [PATCH] Use .get() for discriminator access in generated union encoders The discriminator field may be NotRequired in the TypedDict, so direct key access (x["key"]) triggers pyright's reportTypedDictNotRequiredAccess. Using .get() is safe because a missing key returns None, which won't match any discriminator value and falls through to the next branch. --- src/replit_river/codegen/client.py | 2 +- .../snapshots/test_unknown_enum/enumService/needsEnumObject.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/replit_river/codegen/client.py b/src/replit_river/codegen/client.py index a760acee..977c85a3 100644 --- a/src/replit_river/codegen/client.py +++ b/src/replit_river/codegen/client.py @@ -342,7 +342,7 @@ def flatten_union(tpe: RiverType) -> list[RiverType]: ) typeddict_encoder.append( f""" - if x[{repr(discriminator_name)}] + if x.get({repr(discriminator_name)}) == {repr(discriminator_value)} else """, diff --git a/tests/v1/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py b/tests/v1/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py index 0f0b0d8b..81e98a6d 100644 --- a/tests/v1/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py +++ b/tests/v1/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py @@ -75,7 +75,7 @@ def encode_NeedsenumobjectInput( encode_NeedsenumobjectInputOneOf_in_first( cast("NeedsenumobjectInputOneOf_in_first", x) ) - if x["kind"] == "in_first" + if x.get("kind") == "in_first" else encode_NeedsenumobjectInputOneOf_in_second( cast("NeedsenumobjectInputOneOf_in_second", x) )