Adding and threading through protocol_version

Author: blast-hardcheese

src/replit_river/codegen/client.py

@@ -802,6 +802,7 @@ def render_library_call(
     schema_name: str,
     name: str,
     procedure: RiverProcedure,
+    protocol_version: Literal["v1.1", "v2.0"],
     init_meta: tuple[RiverType, TypeExpression, str] | None,
     input_meta: tuple[RiverType, TypeExpression, str] | None,
     output_meta: tuple[RiverType, TypeExpression, str] | None,
@@ -1022,6 +1023,7 @@ def generate_individual_service(
     schema_name: str,
     schema: RiverService,
     input_base_class: Literal["TypedDict"] | Literal["BaseModel"],
+    protocol_version: Literal["v1.1", "v2.0"],
 ) -> tuple[ModuleName, ClassName, dict[RenderedPath, FileContents]]:
     serdes: list[tuple[list[TypeName], list[ModuleName], list[FileContents]]] = []
 
@@ -1254,6 +1256,7 @@ def combine_or_none(
                 schema_name=schema_name,
                 name=name,
                 procedure=procedure,
+                protocol_version=protocol_version,
                 init_meta=combine_or_none(
                     procedure.init, init_type, render_init_method
                 ),
@@ -1308,6 +1311,7 @@ def generate_river_client_module(
     client_name: str,
     schema_root: RiverSchema,
     typed_dict_inputs: bool,
+    protocol_version: Literal["v1.1", "v2.0"],
 ) -> dict[RenderedPath, FileContents]:
     files: dict[RenderedPath, FileContents] = {}
 
@@ -1332,7 +1336,10 @@ def generate_river_client_module(
     )
     for schema_name, schema in schema_root.services.items():
         module_name, class_name, emitted_files = generate_individual_service(
-            schema_name, schema, input_base_class
+            schema_name,
+            schema,
+            input_base_class,
+            protocol_version,
         )
         files.update(emitted_files)
         modules.append((module_name, class_name))
@@ -1351,12 +1358,16 @@ def schema_to_river_client_codegen(
     client_name: str,
     typed_dict_inputs: bool,
     file_opener: Callable[[Path], TextIO],
+    protocol_version: Literal["v1.1", "v2.0"],
 ) -> None:
     """Generates the lines of a River module."""
     with read_schema() as f:
         schemas = RiverSchemaFile(json.load(f))
     for subpath, contents in generate_river_client_module(
-        client_name, schemas.root, typed_dict_inputs
+        client_name,
+        schemas.root,
+        typed_dict_inputs,
+        protocol_version,
     ).items():
         module_path = Path(target_path).joinpath(subpath)
         module_path.parent.mkdir(mode=0o755, parents=True, exist_ok=True)

src/replit_river/codegen/run.py

@@ -57,11 +57,12 @@ def file_opener(path: Path) -> TextIO:
             return open(path, "w")
 
         schema_to_river_client_codegen(
-            lambda: open(schema_path),
-            target_path,
-            args.client_name,
-            args.typed_dict_inputs,
-            file_opener,
+            read_schema=lambda: open(schema_path),
+            target_path=target_path,
+            client_name=args.client_name,
+            typed_dict_inputs=args.typed_dict_inputs,
+            file_opener=file_opener,
+            protocol_version="v1.1",
         )
     else:
         raise NotImplementedError(f"Unknown command {args.command}")

tests/codegen/snapshot/codegen_snapshot_fixtures.py

@@ -34,6 +34,7 @@ def file_opener(path: Path) -> TextIO:
         client_name=client_name,
         file_opener=file_opener,
         typed_dict_inputs=True,
+        protocol_version="v1.1",
     )
     for path, file in files.items():
         file.seek(0)

tests/codegen/test_rpc.py

@@ -27,11 +27,12 @@ def file_opener(path: Path) -> TextIO:
         return open(path, "w")
 
     schema_to_river_client_codegen(
-        lambda: open("tests/codegen/rpc/schema.json"),
-        "tests/codegen/rpc/generated",
-        "RpcClient",
-        True,
-        file_opener,
+        read_schema=lambda: open("tests/codegen/rpc/schema.json"),
+        target_path="tests/codegen/rpc/generated",
+        client_name="RpcClient",
+        typed_dict_inputs=True,
+        file_opener=file_opener,
+        protocol_version="v1.1",
     )
     importlib.reload(tests.codegen.rpc.generated)