Making input and init types optional, gating on protocol version

Author: blast-hardcheese

src/replit_river/codegen/client.py

@@ -839,8 +839,9 @@ def __init__(self, client: river.Client[Any]):
     for name, procedure in schema.procedures.items():
         module_names = [ModuleName(name)]
         init_type: TypeExpression | None = None
+        init_module_info: list[ModuleName] = []
         if procedure.init:
-            init_type, module_info, init_chunks, encoder_names = encode_type(
+            init_type, init_module_info, init_chunks, encoder_names = encode_type(
                 procedure.init,
                 TypeName(f"{name.title()}Init"),
                 input_base_class,
@@ -850,34 +851,29 @@ def __init__(self, client: river.Client[Any]):
             serdes.append(
                 (
                     [extract_inner_type(init_type), *encoder_names],
-                    module_info,
+                    init_module_info,
                     init_chunks,
                 )
             )
-        input_type, module_info, input_chunks, encoder_names = encode_type(
-            procedure.input,
-            TypeName(f"{name.title()}Input"),
-            input_base_class,
-            module_names,
-            permit_unknown_members=False,
-        )
-        input_type_name = extract_inner_type(input_type)
-        input_type_type_adapter_name = TypeName(
-            f"{render_literal_type(input_type_name)}TypeAdapter"
-        )
-        serdes.append(
-            (
-                [extract_inner_type(input_type), *encoder_names],
-                module_info,
-                input_chunks,
+        input_type: TypeExpression | None = None
+        input_module_info: list[ModuleName] = []
+        if procedure.input:
+            input_type, input_module_info, input_chunks, encoder_names = encode_type(
+                procedure.input,
+                TypeName(f"{name.title()}Input"),
+                input_base_class,
+                module_names,
+                permit_unknown_members=False,
             )
-        )
-        serdes.append(
-            _type_adapter_definition(
-                input_type_type_adapter_name, input_type, module_info
+            serdes.append(
+                (
+                    [extract_inner_type(input_type), *encoder_names],
+                    input_module_info,
+                    input_chunks,
+                )
             )
-        )
-        output_type, module_info, output_chunks, encoder_names = encode_type(
+
+        output_type, output_module_info, output_chunks, encoder_names = encode_type(
             procedure.output,
             TypeName(f"{name.title()}Output"),
             "BaseModel",
@@ -888,7 +884,7 @@ def __init__(self, client: river.Client[Any]):
         serdes.append(
             (
                 [output_type_name, *encoder_names],
-                module_info,
+                output_module_info,
                 output_chunks,
             )
         )
@@ -897,12 +893,12 @@ def __init__(self, client: river.Client[Any]):
         )
         serdes.append(
             _type_adapter_definition(
-                output_type_type_adapter_name, output_type, module_info
+                output_type_type_adapter_name, output_type, output_module_info
             )
         )
-        output_module_info = module_info
+
         if procedure.errors:
-            error_type, module_info, errors_chunks, encoder_names = encode_type(
+            error_type, error_module_info, errors_chunks, encoder_names = encode_type(
                 procedure.errors,
                 TypeName(f"{name.title()}Errors"),
                 "RiverError",
@@ -914,7 +910,7 @@ def __init__(self, client: river.Client[Any]):
                 error_type = error_type_name
             else:
                 error_type_name = extract_inner_type(error_type)
-                serdes.append(([error_type_name], module_info, errors_chunks))
+                serdes.append(([error_type_name], error_module_info, errors_chunks))
 
         else:
             error_type_name = TypeName("RiverError")
@@ -924,11 +920,9 @@ def __init__(self, client: river.Client[Any]):
             f"{render_literal_type(error_type_name)}TypeAdapter"
         )
         if error_type_type_adapter_name.value != "RiverErrorTypeAdapter":
-            if len(module_info) == 0:
-                module_info = output_module_info
             serdes.append(
                 _type_adapter_definition(
-                    error_type_type_adapter_name, error_type, module_info
+                    error_type_type_adapter_name, error_type, output_module_info
                 )
             )
         output_or_error_type = UnionTypeExpr([output_type, error_type_name])
@@ -975,48 +969,50 @@ def __init__(self, client: river.Client[Any]):
                 )
                 serdes.append(
                     _type_adapter_definition(
-                        init_type_type_adapter_name, init_type, module_info
+                        init_type_type_adapter_name, init_type, init_module_info
                     )
                 )
                 render_init_method = f"""\
                     lambda x: {render_type_expr(init_type_type_adapter_name)}
                         .validate_python
                 """
 
-        assert init_type is None or render_init_method, (
-            f"Unable to derive the init encoder from: {init_type}"
-        )
-
         # Input renderer
         render_input_method: str | None = None
-        if input_base_class == "TypedDict":
-            if is_literal(procedure.input):
-                render_input_method = "lambda x: x"
-            elif isinstance(
-                procedure.input, RiverConcreteType
-            ) and procedure.input.type in ["array"]:
-                match input_type:
-                    case ListTypeExpr(list_type):
-                        render_input_method = f"""\
-                        lambda xs: [
-                            encode_{render_literal_type(list_type)}(x) for x in xs
-                        ]
-                        """
-            else:
-                render_input_method = f"encode_{render_literal_type(input_type)}"
-        else:
-            render_input_method = f"""\
-                            lambda x: {render_type_expr(input_type_type_adapter_name)}
-                              .dump_python(
-                                x, # type: ignore[arg-type]
-                                by_alias=True,
-                                exclude_none=True,
-                              )
+        if input_type and procedure.input is not None:
+            if input_base_class == "TypedDict":
+                if is_literal(procedure.input):
+                    render_input_method = "lambda x: x"
+                elif isinstance(
+                    procedure.input, RiverConcreteType
+                ) and procedure.input.type in ["array"]:
+                    match input_type:
+                        case ListTypeExpr(list_type):
+                            render_input_method = f"""\
+                            lambda xs: [
+                                encode_{render_literal_type(list_type)}(x) for x in xs
+                            ]
                             """
-
-        assert render_input_method, (
-            f"Unable to derive the input encoder from: {input_type}"
-        )
+                else:
+                    render_input_method = f"encode_{render_literal_type(input_type)}"
+            else:
+                input_type_name = extract_inner_type(input_type)
+                input_type_type_adapter_name = TypeName(
+                    f"{render_literal_type(input_type_name)}TypeAdapter"
+                )
+                serdes.append(
+                    _type_adapter_definition(
+                        input_type_type_adapter_name, input_type, input_module_info
+                    )
+                )
+                render_input_method = f"""\
+                                lambda x: {render_type_expr(input_type_type_adapter_name)}
+                                  .dump_python(
+                                    x, # type: ignore[arg-type]
+                                    by_alias=True,
+                                    exclude_none=True,
+                                  )
+                                """
 
         if isinstance(output_type, NoneTypeExpr):
             parse_output_method = "lambda x: None"
@@ -1075,7 +1071,6 @@ async def {name}(
             )
         elif procedure.type == "upload":
             if init_type:
-                assert render_init_method, "Expected an init renderer!"
                 current_chunks.extend(
                     [
                         reindent(
@@ -1132,7 +1127,6 @@ async def {name}(
                 ]
             )
             if init_type:
-                assert render_init_method, "Expected an init renderer!"
                 current_chunks.extend(
                     [
                         reindent(