Fix mypy arg-type errors in generated discriminated union encoders

src/replit_river/codegen/client.py

@@ -80,6 +80,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 
@@ -301,9 +302,10 @@ def flatten_union(tpe: RiverType) -> list[RiverType]:
                                 # "encoder_names" is only a TypedDict thing
                                 encoder_names.add(encoder_name)
                             _field_name = render_literal_type(encoder_name)
+                            _type_name = render_literal_type(type_name)
                             typeddict_encoder.append(
                                 f"""\
-                                {_field_name}(x) # type: ignore[arg-type]
+                                {_field_name}(cast('{_type_name}', x))
                                 """.strip()
                             )
                             if local_discriminators:
@@ -333,8 +335,10 @@ def flatten_union(tpe: RiverType) -> list[RiverType]:
                         # TODO(dstewart): Figure out why uncommenting this breaks
                         #                 generated code
                         # encoder_names.add(encoder_name)
+                        _type_name = render_literal_type(type_name)
                         typeddict_encoder.append(
-                            f"{render_literal_type(encoder_name)}(x)"
+                            f"{render_literal_type(encoder_name)}"
+                            f"(cast('{_type_name}', x))"
                         )
                     typeddict_encoder.append(
                         f"""

tests/v1/codegen/rpc/generated/test_service/rpc_method.py

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 

tests/v1/codegen/rpc/generated_special_chars/test_service/rpc_method.py

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 

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

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 

tests/v1/codegen/snapshot/snapshots/test_basic_stream/test_service/emit_error.py

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 

tests/v1/codegen/snapshot/snapshots/test_basic_stream/test_service/stream_method.py

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 

tests/v1/codegen/snapshot/snapshots/test_pathological_types/test_service/pathological_method.py

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 

tests/v1/codegen/snapshot/snapshots/test_recursive_types/recursiveService/getTree.py

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 

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

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 

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

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 
@@ -71,9 +72,13 @@ def encode_NeedsenumobjectInput(
     x: "NeedsenumobjectInput",
 ) -> Any:
     return (
-        encode_NeedsenumobjectInputOneOf_in_first(x)
+        encode_NeedsenumobjectInputOneOf_in_first(
+            cast("NeedsenumobjectInputOneOf_in_first", x)
+        )
         if x["kind"] == "in_first"
-        else encode_NeedsenumobjectInputOneOf_in_second(x)
+        else encode_NeedsenumobjectInputOneOf_in_second(
+            cast("NeedsenumobjectInputOneOf_in_second", x)
+        )
     )
 
 

tests/v2/codegen/snapshot/snapshots/test_basic_rpc/test_service/rpc_method.py

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated
 

tests/v2/codegen/snapshot/snapshots/test_basic_stream/test_service/stream_method.py

@@ -7,6 +7,7 @@
     Mapping,
     NotRequired,
     TypedDict,
+    cast,
 )
 from typing_extensions import Annotated