Modularizing client codegen

Author: blast-hardcheese

replit_river/codegen/client.py

@@ -1,11 +1,13 @@
 import json
 import re
+from pathlib import Path
 from textwrap import dedent, indent
 from typing import (
     Any,
     Dict,
     List,
     Literal,
+    NewType,
     Optional,
     OrderedDict,
     Sequence,
@@ -18,6 +20,11 @@
 import black
 from pydantic import BaseModel, Field, RootModel
 
+ModuleName = NewType("ModuleName", str)
+ClassName = NewType("ClassName", str)
+FileContents = NewType("FileContents", str)
+HandshakeType = NewType("HandshakeType", str)
+
 _NON_ALNUM_RE = re.compile(r"[^a-zA-Z0-9_]+")
 _LITERAL_RE = re.compile(r"^Literal\[(.+)\]$")
 
@@ -485,13 +492,50 @@ def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
     return (prefix, chunks)
 
 
-def generate_individual_service(schema_name: str, handshake_type: str, schema: RiverService, input_base_class: Literal['TypedDict'] | Literal['BaseModel']) -> list[str]:
+def generate_common_client(
+    client_name: str,
+    handshake_type: HandshakeType,
+    handshake_chunks: Sequence[str],
+    modules: list[Tuple[ModuleName, ClassName]],
+) -> FileContents:
+    chunks: list[str] = [FILE_HEADER]
+    chunks.extend(
+        [
+            f"from .{model_name} import {class_name}"
+            for model_name, class_name in modules
+        ]
+    )
+    chunks.extend(handshake_chunks)
+    chunks.extend(
+        [
+            dedent(
+                f"""\
+                class {client_name}:
+                  def __init__(self, client: river.Client[{handshake_type}]):
+                """.rstrip()
+            )
+        ]
+    )
+    for module_name, class_name in modules:
+        chunks.append(
+            f"    self.{module_name} = {class_name}(client)",
+        )
+
+    return FileContents("\n".join(chunks))
+
+
+def generate_individual_service(
+    schema_name: str,
+    schema: RiverService,
+    input_base_class: Literal["TypedDict"] | Literal["BaseModel"],
+) -> Tuple[Path, ModuleName, ClassName, FileContents]:
     serdes: list[str] = []
+    class_name = ClassName(f"{schema_name.title()}Service")
     current_chunks: List[str] = [
         dedent(
             f"""\
-              class {schema_name.title()}Service:
-                def __init__(self, client: river.Client[{handshake_type}]):
+              class {class_name}:
+                def __init__(self, client: river.Client[Any]):
                   self.client = client
             """
         ),
@@ -549,15 +593,13 @@ def __init__(self, client: river.Client[{handshake_type}]):
             """.rstrip()
 
         # Init renderer
-        if input_base_class == 'TypedDict' and init_type:
+        if input_base_class == "TypedDict" and init_type:
             if is_literal(procedure.input):
                 render_init_method = "lambda x: x"
             elif isinstance(
                 procedure.input, RiverConcreteType
             ) and procedure.input.type in ["array"]:
-                assert init_type.startswith(
-                    "List["
-                )  # in case we change to list[...]
+                assert init_type.startswith("List[")  # in case we change to list[...]
                 _init_type_name = init_type[len("List[") : -len("]")]
                 render_init_method = (
                     f"lambda xs: [encode_{_init_type_name}(x) for x in xs]"
@@ -575,15 +617,13 @@ def __init__(self, client: river.Client[{handshake_type}]):
             render_init_method = "lambda x: x"
 
         # Input renderer
-        if input_base_class == 'TypedDict':
+        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"]:
-                assert input_type.startswith(
-                    "List["
-                )  # in case we change to list[...]
+                assert input_type.startswith("List[")  # in case we change to list[...]
                 _input_type_name = input_type[len("List[") : -len("]")]
                 render_input_method = (
                     f"lambda xs: [encode_{_input_type_name}(x) for x in xs]"
@@ -760,45 +800,48 @@ async def {name}(
                 )
 
         current_chunks.append("")
-    return current_chunks
+    return (
+        Path(f"{schema_name}.py"),
+        ModuleName(schema_name),
+        class_name,
+        FileContents("\n".join([FILE_HEADER] + serdes + current_chunks)),
+    )
 
 
 def generate_river_client_module(
     client_name: str,
     schema_root: RiverSchema,
     typed_dict_inputs: bool,
-) -> Sequence[str]:
-    chunks: List[str] = [FILE_HEADER]
+) -> dict[Path, FileContents]:
+    files: dict[Path, FileContents] = {}
 
+    # Negotiate handshake shape
+    handshake_chunks: Sequence[str] = []
     if schema_root.handshakeSchema is not None:
-        (handshake_type, handshake_chunks) = encode_type(
+        (_handshake_type, handshake_chunks) = encode_type(
             schema_root.handshakeSchema, "HandshakeSchema", "BaseModel"
         )
-        chunks.extend(handshake_chunks)
+        handshake_type = HandshakeType(_handshake_type)
     else:
-        handshake_type = "Literal[None]"
+        handshake_type = HandshakeType("Literal[None]")
 
-    input_base_class: Literal['TypedDict'] | Literal['BaseModel'] = "TypedDict" if typed_dict_inputs else "BaseModel"
-    for schema_name, schema in schema_root.services.items():
-        current_chunks = generate_individual_service(schema_name, handshake_type, schema, input_base_class)
-        chunks.extend(current_chunks)
-
-    chunks.extend(
-        [
-            dedent(
-                f"""\
-                class {client_name}:
-                  def __init__(self, client: river.Client[{handshake_type}]):
-                """.rstrip()
-            )
-        ]
+    modules: list[Tuple[ModuleName, ClassName]] = []
+    input_base_class: Literal["TypedDict"] | Literal["BaseModel"] = (
+        "TypedDict" if typed_dict_inputs else "BaseModel"
     )
     for schema_name, schema in schema_root.services.items():
-        chunks.append(
-            f"    self.{schema_name} = {schema_name.title()}Service(client)",
+        path, module_name, class_name, file_contents = generate_individual_service(
+            schema_name, schema, input_base_class
         )
+        files[path] = file_contents
+        modules.append((module_name, class_name))
+
+    main_contents = generate_common_client(
+        client_name, handshake_type, handshake_chunks, modules
+    )
+    files[Path("__init__.py")] = main_contents
 
-    return chunks
+    return files
 
 
 def schema_to_river_client_codegen(
@@ -810,14 +853,18 @@ def schema_to_river_client_codegen(
     """Generates the lines of a River module."""
     with open(schema_path) as f:
         schemas = RiverSchemaFile(json.load(f))
-    with open(target_path, "w") as f:
-        s = "\n".join(
-            generate_river_client_module(client_name, schemas.root, typed_dict_inputs)
-        )
-        try:
-            f.write(
-                black.format_str(s, mode=black.FileMode(string_normalization=False))
-            )
-        except:
-            f.write(s)
-            raise
+    for subpath, contents in generate_river_client_module(
+        client_name, schemas.root, typed_dict_inputs
+    ).items():
+        module_path = Path(target_path).joinpath(subpath)
+        module_path.parent.mkdir(mode=0o755, parents=True, exist_ok=True)
+        with open(module_path, "w") as f:
+            try:
+                f.write(
+                    black.format_str(
+                        contents, mode=black.FileMode(string_normalization=False)
+                    )
+                )
+            except:
+                f.write(contents)
+                raise