Permit method filtering based on supplied file

Author: blast-hardcheese

src/replit_river/codegen/client.py

@@ -1113,6 +1113,7 @@ def generate_individual_service(
     schema: RiverService,
     input_base_class: Literal["TypedDict"] | Literal["BaseModel"],
     protocol_version: Literal["v1.1", "v2.0"],
+    method_filter: set[str] | None,
 ) -> tuple[ModuleName, ClassName, dict[RenderedPath, FileContents]]:
     serdes: list[tuple[list[TypeName], list[ModuleName], list[FileContents]]] = []
 
@@ -1155,6 +1156,8 @@ def __init__(self, client: {client_module}.Client[Any]):
         ),
     ]
     for name, procedure in schema.procedures.items():
+        if method_filter and (schema_name + "." + name) in method_filter:
+            continue
         module_names = [ModuleName(name)]
         init_type: TypeExpression | None = None
         init_module_info: list[ModuleName] = []
@@ -1408,6 +1411,7 @@ def generate_river_client_module(
     schema_root: RiverSchema,
     typed_dict_inputs: bool,
     protocol_version: Literal["v1.1", "v2.0"],
+    method_filter: set[str] | None,
 ) -> dict[RenderedPath, FileContents]:
     files: dict[RenderedPath, FileContents] = {}
 
@@ -1436,9 +1440,12 @@ def generate_river_client_module(
             schema,
             input_base_class,
             protocol_version,
+            method_filter,
         )
-        files.update(emitted_files)
-        modules.append((module_name, class_name))
+        if emitted_files:
+            # Short-cut if we didn't actually emit anything
+            files.update(emitted_files)
+            modules.append((module_name, class_name))
 
     main_contents = generate_common_client(
         client_name,
@@ -1459,6 +1466,7 @@ def schema_to_river_client_codegen(
     typed_dict_inputs: bool,
     file_opener: Callable[[Path], TextIO],
     protocol_version: Literal["v1.1", "v2.0"],
+    method_filter: set[str] | None,
 ) -> None:
     """Generates the lines of a River module."""
     with read_schema() as f:
@@ -1468,6 +1476,7 @@ def schema_to_river_client_codegen(
         schemas.root,
         typed_dict_inputs,
         protocol_version,
+        method_filter,
     ).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

@@ -1,5 +1,6 @@
 import argparse
 import os.path
+import pathlib
 from pathlib import Path
 from typing import TextIO
 
@@ -45,9 +46,20 @@ def main() -> None:
         default="v1.1",
         choices=["v1.1", "v2.0"],
     )
+    client.add_argument(
+        "--method-filter",
+        help="Only generate a subset of the specified methods",
+        action="store",
+        type=pathlib.Path,
+    )
     client.add_argument("schema", help="schema file")
     args = parser.parse_args()
 
+    method_filter: set[str] | None = None
+    if args.method_filter:
+        with open(args.method_filter) as handle:
+            method_filter = set(x.strip() for x in handle.readlines())
+
     if args.command == "server":
         proto_path = os.path.abspath(args.proto)
         target_directory = os.path.abspath(args.output)
@@ -70,6 +82,7 @@ def file_opener(path: Path) -> TextIO:
             typed_dict_inputs=args.typed_dict_inputs,
             file_opener=file_opener,
             protocol_version=args.protocol_version,
+            method_filter=method_filter,
         )
     else:
         raise NotImplementedError(f"Unknown command {args.command}")

tests/codegen/snapshot/codegen_snapshot_fixtures.py

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

tests/codegen/test_rpc.py

@@ -33,6 +33,7 @@ def file_opener(path: Path) -> TextIO:
         typed_dict_inputs=True,
         file_opener=file_opener,
         protocol_version="v1.1",
+        method_filter=None,
     )
     importlib.reload(tests.codegen.rpc.generated)