client.py

import json
import re
import subprocess
from pathlib import Path
from textwrap import dedent
from typing import (
    Any,
    Callable,
    Literal,
    OrderedDict,
    Sequence,
    Set,
    TextIO,
    cast,
)

from pydantic import BaseModel, Field, RootModel

from replit_river.codegen.format import reindent
from replit_river.codegen.typing import (
    ClassName,
    DictTypeExpr,
    FileContents,
    HandshakeType,
    ListTypeExpr,
    LiteralTypeExpr,
    ModuleName,
    NoneTypeExpr,
    OpenUnionTypeExpr,
    RenderedPath,
    TypeExpression,
    TypeName,
    UnionTypeExpr,
    extract_inner_type,
    render_literal_type,
    render_type_expr,
)

_NON_ALNUM_RE = re.compile(r"[^a-zA-Z0-9_]+")

# Literal is here because HandshakeType can be Literal[None]
ROOT_FILE_HEADER = dedent(
    """\
# Code generated by river.codegen. DO NOT EDIT.
from pydantic import BaseModel
from typing import Literal

import replit_river as river

"""
)

SERVICE_FILE_HEADER = dedent(
    """\
# Code generated by river.codegen. DO NOT EDIT.
from collections.abc import AsyncIterable, AsyncIterator
from typing import Any
import datetime

from pydantic import TypeAdapter

from replit_river.error_schema import RiverError, RiverErrorTypeAdapter
import replit_river as river

"""
)

FILE_HEADER = dedent(
    """\
# Code generated by river.codegen. DO NOT EDIT.
from collections.abc import AsyncIterable, AsyncIterator
import datetime
from typing import (
    Any,
    Literal,
    Mapping,
    NotRequired,
    TypedDict,
)
from typing_extensions import Annotated

from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
from replit_river.error_schema import RiverError
from replit_river.client import RiverUnknownValue, translate_unknown_value

import replit_river as river

"""
)


class RiverConcreteType(BaseModel):
    type: str | None = Field(default=None)
    properties: dict[str, "RiverType"] = Field(default_factory=lambda: dict())
    required: Set[str] = Field(default=set())
    items: "RiverType | None" = Field(default=None)
    const: str | int | None = Field(default=None)
    patternProperties: dict[str, "RiverType"] = Field(default_factory=lambda: dict())


class RiverUnionType(BaseModel):
    anyOf: list["RiverType"]


class RiverIntersectionType(BaseModel):
    allOf: list["RiverType"]


class RiverNotType(BaseModel):
    """This is used to represent void / never."""

    not_: Any = Field(..., alias="not")


RiverType = RiverConcreteType | RiverUnionType | RiverNotType | RiverIntersectionType


class RiverProcedure(BaseModel):
    init: RiverType | None = Field(default=None)
    input: RiverType
    output: RiverType
    errors: RiverType | None = Field(default=None)
    type: (
        Literal["rpc"] | Literal["stream"] | Literal["subscription"] | Literal["upload"]
    )
    description: str | None = Field(default=None)


class RiverService(BaseModel):
    procedures: dict[str, RiverProcedure]


class RiverSchema(BaseModel):
    services: dict[str, RiverService]
    handshakeSchema: RiverConcreteType | None = Field(default=None)


RiverSchemaFile = RootModel[RiverSchema]


def is_literal(tpe: RiverType) -> bool:
    if isinstance(tpe, RiverUnionType):
        return all(is_literal(t) for t in tpe.anyOf)
    elif isinstance(tpe, RiverConcreteType):
        return tpe.type not in set(["object", "array"])
    else:
        return False


def encode_type(
    type: RiverType,
    prefix: TypeName,
    base_model: str,
    in_module: list[ModuleName],
    permit_unknown_members: bool,
) -> tuple[TypeExpression, list[ModuleName], list[FileContents], set[TypeName]]:
    encoder_name: TypeName | None = None  # defining this up here to placate mypy
    chunks: list[FileContents] = []
    if isinstance(type, RiverNotType):
        return (NoneTypeExpr(), [], [], set())
    elif isinstance(type, RiverUnionType):
        typeddict_encoder = list[str]()
        encoder_names: set[TypeName] = set()

        # First check if it's a discriminated union. Typebox currently doesn't have
        # a way of expressing the intention of having a discriminated union. So we
        # do a bit of detection if that is structurally true by checking that all the
        # types in the anyOf are objects, have properties, and have one property common
        # to all the alternatives that has a literal value.
        def flatten_union(tpe: RiverType) -> list[RiverType]:
            if isinstance(tpe, RiverUnionType):
                return [u for t in tpe.anyOf for u in flatten_union(t)]
            else:
                return [tpe]

        original_type = type

        type = RiverUnionType(anyOf=flatten_union(type))

        one_of_candidate_types: list[RiverConcreteType] = [
            t
            for _t in type.anyOf
            for t in (_t.anyOf if isinstance(_t, RiverUnionType) else [_t])
            if isinstance(t, RiverConcreteType)
            and t.type == "object"
            and t.properties
            and (not t.patternProperties or "^(.*)$" not in t.patternProperties)
        ]
        if len(type.anyOf) > 0 and len(type.anyOf) == len(one_of_candidate_types):
            # We have established that it is a union-of-objects. Now let's see if
            # there is a discriminator field common among all options.
            literal_fields = set[str]()
            for i, oneof_t in enumerate(one_of_candidate_types):
                lf = set[str](
                    name
                    for name, prop in oneof_t.properties.items()
                    if isinstance(prop, RiverConcreteType)
                    and prop.type in ("string", "number", "boolean")
                    and prop.const is not None
                )
                if i == 0:
                    literal_fields = lf
                else:
                    literal_fields.intersection_update(lf)
                if not literal_fields:
                    # There are no more candidates.
                    break
            if len(literal_fields) == 1:
                # Hooray! we found a discriminated union.
                discriminator_name = literal_fields.pop()
                one_of_pending = OrderedDict[str, tuple[str, list[RiverConcreteType]]]()

                for oneof_t in one_of_candidate_types:
                    discriminator_value = [
                        _NON_ALNUM_RE.sub("", str(prop.const))
                        for name, prop in sorted(
                            list(oneof_t.properties.items()), key=lambda kv: kv[0]
                        )
                        if isinstance(prop, RiverConcreteType)
                        and name == discriminator_name
                        and prop.const is not None
                    ].pop()
                    one_of_pending.setdefault(
                        f"{render_literal_type(prefix)}OneOf_{discriminator_value}",
                        (discriminator_value, []),
                    )[1].append(oneof_t)

                one_of: list[TypeExpression] = []
                if discriminator_name == "$kind":
                    discriminator_name = "kind"
                for pfx, (discriminator_value, oneof_ts) in one_of_pending.items():
                    if len(oneof_ts) > 1:
                        typeddict_encoder.append("(")
                        # Tricky bit. We need to derive our own discriminator based
                        # on known members. Be careful.

                        common_members = set[str]()
                        for i, oneof_t in enumerate(oneof_ts):
                            if i == 0:
                                common_members = set(oneof_t.properties.keys())
                            else:
                                common_members.intersection_update(
                                    oneof_t.properties.keys()
                                )

                        for i, oneof_t in enumerate(oneof_ts):
                            type_name, _, contents, _ = encode_type(
                                oneof_ts[i],
                                TypeName(f"{pfx}{i}"),
                                base_model,
                                in_module,
                                permit_unknown_members=permit_unknown_members,
                            )
                            one_of.append(type_name)
                            chunks.extend(contents)
                            local_discriminators = set(
                                oneof_t.properties.keys()
                            ).difference(common_members)
                            encoder_name = TypeName(
                                f"encode_{render_literal_type(type_name)}"
                            )
                            encoder_names.add(encoder_name)
                            _field_name = render_literal_type(encoder_name)
                            typeddict_encoder.append(
                                f"""\
                                {_field_name}(x) # type: ignore[arg-type]
                                """.strip()
                            )
                            if local_discriminators:
                                local_discriminator = sorted(local_discriminators).pop()
                            else:
                                local_discriminator = "FIXME: Ambiguous discriminators"
                            typeddict_encoder.append(
                                f" if {repr(local_discriminator)} in x else "
                            )
                        typeddict_encoder.pop()  # Drop the last ternary
                        typeddict_encoder.append(")")
                    else:
                        oneof_t = oneof_ts[0]
                        type_name, _, contents, _ = encode_type(
                            oneof_t,
                            TypeName(pfx),
                            base_model,
                            in_module,
                            permit_unknown_members=permit_unknown_members,
                        )
                        one_of.append(type_name)
                        chunks.extend(contents)
                        encoder_name = TypeName(
                            f"encode_{render_literal_type(type_name)}"
                        )
                        # TODO(dstewart): Figure out why uncommenting this breaks
                        #                 generated code
                        # encoder_names.add(encoder_name)
                        typeddict_encoder.append(
                            f"{render_literal_type(encoder_name)}(x)"
                        )
                    typeddict_encoder.append(
                        f"""
                            if x[{repr(discriminator_name)}]
                            == {repr(discriminator_value)}
                            else
                        """,
                    )
                union: TypeExpression
                if permit_unknown_members:
                    union = OpenUnionTypeExpr(UnionTypeExpr(one_of))
                else:
                    union = UnionTypeExpr(one_of)
                chunks.append(
                    FileContents(
                        f"{render_literal_type(prefix)} = {render_type_expr(union)}"
                    )
                )
                chunks.append(FileContents(""))

                if base_model == "TypedDict":
                    encoder_name = TypeName(f"encode_{render_literal_type(prefix)}")
                    encoder_names.add(encoder_name)
                    _field_name = render_literal_type(encoder_name)
                    typeddict_encoder = typeddict_encoder[:-1]  # Drop the last ternary
                    chunks.append(
                        FileContents(
                            "\n".join(
                                [
                                    dedent(
                                        f"""\
                            def {_field_name}(
                                x: {repr(render_literal_type(prefix))},
                            ) -> Any:
                                return (
                                    {
                                            reindent(
                                                "                                    ",
                                                "\n".join(typeddict_encoder),
                                            )
                                        }
                                )
                                        """.rstrip()
                                    )
                                ]
                            )
                        )
                    )
                return (prefix, in_module, chunks, encoder_names)
            # End of stable union detection
        # Restore the non-flattened union type
        type = original_type
        any_of: list[TypeExpression] = []

        typeddict_encoder = []
        for i, t in enumerate(type.anyOf):
            type_name, _, contents, _ = encode_type(
                t,
                TypeName(f"{render_literal_type(prefix)}AnyOf_{i}"),
                base_model,
                in_module,
                permit_unknown_members=permit_unknown_members,
            )
            any_of.append(type_name)
            chunks.extend(contents)
            if isinstance(t, RiverConcreteType):
                if t.type == "string":
                    typeddict_encoder.extend(["x", " if isinstance(x, str) else "])
                else:
                    # TODO(dstewart): This structure changed since we were incorrectly
                    #                 leaking ListTypeExprs into codegen. This generated
                    #                 code is probably wrong.
                    match type_name:
                        case ListTypeExpr(inner_type_name):
                            typeddict_encoder.append(
                                f"encode_{render_literal_type(inner_type_name)}(x)"
                            )
                        case LiteralTypeExpr(const):
                            typeddict_encoder.append(repr(const))
                        case TypeName(value):
                            typeddict_encoder.append(f"encode_{value}(x)")
                        case NoneTypeExpr():
                            typeddict_encoder.append("None")
                        case other:
                            _o2: DictTypeExpr | OpenUnionTypeExpr | UnionTypeExpr = (
                                other
                            )
                            raise ValueError(f"What does it mean to have {_o2} here?")
        if permit_unknown_members:
            union = OpenUnionTypeExpr(UnionTypeExpr(any_of))
        else:
            union = UnionTypeExpr(any_of)
        if is_literal(type):
            typeddict_encoder = ["x"]
        chunks.append(
            FileContents(f"{render_literal_type(prefix)} = {render_type_expr(union)}")
        )
        if base_model == "TypedDict":
            encoder_name = TypeName(f"encode_{render_literal_type(prefix)}")
            encoder_names.add(encoder_name)
            _field_name = render_literal_type(encoder_name)
            chunks.append(
                FileContents(
                    dedent(
                        f"""
                    def {_field_name}(x: {repr(render_literal_type(prefix))}) -> Any:
                        return (
                            {
                            reindent(
                                "                            ",
                                "\n".join(typeddict_encoder),
                            )
                        }
                        )
                        """
                    )
                )
            )
        return (prefix, in_module, chunks, encoder_names)
    elif isinstance(type, RiverIntersectionType):

        def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
            if isinstance(tpe, RiverUnionType):
                return [t for p in tpe.anyOf for t in extract_props(p)]
            elif isinstance(tpe, RiverConcreteType):
                return [tpe.properties]
            elif isinstance(tpe, RiverIntersectionType):
                return [t for p in tpe.allOf for t in extract_props(p)]
            else:
                return []

        combined = {}
        for x in extract_props(type):
            combined.update(x)
        return encode_type(
            RiverConcreteType(type="object", properties=combined),
            prefix,
            base_model,
            in_module,
            permit_unknown_members=permit_unknown_members,
        )
    elif isinstance(type, RiverConcreteType):
        typeddict_encoder = list[str]()
        if type.type is None:
            # Handle the case where type is not specified
            typeddict_encoder.append("x")
            return (TypeName("Any"), [], [], set())
        elif type.type == "string":
            if type.const:
                typeddict_encoder.append(repr(type.const))
                return (LiteralTypeExpr(type.const), [], [], set())
            else:
                typeddict_encoder.append("x")
                return (TypeName("str"), [], [], set())
        elif type.type == "Uint8Array":
            typeddict_encoder.append("x.decode()")
            return (TypeName("bytes"), [], [], set())
        elif type.type == "number":
            if type.const is not None:
                # enums are represented as const number in the schema
                typeddict_encoder.append(f"{type.const}")
                return (LiteralTypeExpr(type.const), [], [], set())
            typeddict_encoder.append("x")
            return (TypeName("float"), [], [], set())
        elif type.type == "integer":
            if type.const is not None:
                # enums are represented as const number in the schema
                typeddict_encoder.append(f"{type.const}")
                return (LiteralTypeExpr(type.const), [], [], set())
            typeddict_encoder.append("x")
            return (TypeName("int"), [], [], set())
        elif type.type == "boolean":
            typeddict_encoder.append("x")
            return (TypeName("bool"), [], [], set())
        elif type.type == "null" or type.type == "undefined":
            typeddict_encoder.append("None")
            return (NoneTypeExpr(), [], [], set())
        elif type.type == "Date":
            typeddict_encoder.append("TODO: dstewart")
            return (TypeName("datetime.datetime"), [], [], set())
        elif type.type == "array" and type.items:
            type_name, module_info, type_chunks, encoder_names = encode_type(
                type.items,
                prefix,
                base_model,
                in_module,
                permit_unknown_members=permit_unknown_members,
            )
            typeddict_encoder.append("TODO: dstewart")
            return (ListTypeExpr(type_name), module_info, type_chunks, encoder_names)
        elif (
            type.type == "object"
            and type.patternProperties
            and "^(.*)$" in type.patternProperties
        ):
            type_name, module_info, type_chunks, encoder_names = encode_type(
                type.patternProperties["^(.*)$"],
                prefix,
                base_model,
                in_module,
                permit_unknown_members=permit_unknown_members,
            )
            # TODO(dstewart): This structure changed since we were incorrectly leaking
            #                 ListTypeExprs into codegen. This generated code is
            #                 probably wrong.
            match type_name:
                case ListTypeExpr(inner_type_name):
                    typeddict_encoder.append(
                        f"encode_{render_literal_type(inner_type_name)}(x)"
                    )
                case DictTypeExpr(_):
                    raise ValueError(
                        "What does it mean to try and encode a dict in this position?"
                    )
                case LiteralTypeExpr(const):
                    typeddict_encoder.append(repr(const))
                case TypeName(value):
                    typeddict_encoder.append(f"encode_{value}(x)")
                case other:
                    _o1: NoneTypeExpr | OpenUnionTypeExpr | UnionTypeExpr = other
                    raise ValueError(f"What does it mean to have {_o1} here?")
            return (DictTypeExpr(type_name), module_info, type_chunks, encoder_names)
        assert type.type == "object", type.type

        current_chunks: list[str] = [
            f"class {render_literal_type(prefix)}({base_model}):"
        ]
        # For the encoder path, do we need "x" to be bound?
        # lambda x: ... vs lambda _: {}
        needs_binding = False
        encoder_names = set()
        if type.properties:
            needs_binding = True
            typeddict_encoder.append("{")
            for (
                name,
                prop,
            ) in sorted(list(type.properties.items()), key=lambda xs: xs[0]):
                typeddict_encoder.append(f"{repr(name)}:")
                type_name, _, contents, _ = encode_type(
                    prop,
                    TypeName(prefix.value + name.title()),
                    base_model,
                    in_module,
                    permit_unknown_members=permit_unknown_members,
                )
                encoder_name = None
                chunks.extend(contents)
                if base_model == "TypedDict":
                    if isinstance(prop, RiverNotType):
                        typeddict_encoder.append("'not implemented'")
                    elif isinstance(prop, RiverUnionType):
                        encoder_name = TypeName(
                            f"encode_{render_literal_type(type_name)}"
                        )
                        encoder_names.add(encoder_name)
                        typeddict_encoder.append(
                            f"{render_literal_type(encoder_name)}(x[{repr(name)}])"
                        )
                        if name not in type.required:
                            typeddict_encoder.append(
                                f"if {repr(name)} in x and x[{repr(name)}] else None"
                            )
                    elif isinstance(prop, RiverIntersectionType):
                        encoder_name = TypeName(
                            f"encode_{render_literal_type(type_name)}"
                        )
                        encoder_names.add(encoder_name)
                        typeddict_encoder.append(f"{encoder_name}(x[{repr(name)}])")
                    elif isinstance(prop, RiverConcreteType):
                        if name == "$kind":
                            safe_name = "kind"
                        else:
                            safe_name = name
                        if prop.type == "object" and not prop.patternProperties:
                            encoder_name = TypeName(
                                f"encode_{render_literal_type(type_name)}"
                            )
                            encoder_names.add(encoder_name)
                            typeddict_encoder.append(
                                f"{render_literal_type(encoder_name)}(x[{repr(safe_name)}])"
                            )
                            if name not in prop.required:
                                typeddict_encoder.append(
                                    dedent(
                                        f"""
                                        if {repr(safe_name)} in x
                                        and x[{repr(safe_name)}] is not None
                                        else None
                                    """
                                    )
                                )
                        elif prop.type == "array":
                            items = cast(RiverConcreteType, prop).items
                            assert items, "Somehow items was none"
                            if is_literal(cast(RiverType, items)):
                                if name in prop.required:
                                    typeddict_encoder.append(f"x[{repr(name)}]")
                                else:
                                    typeddict_encoder.append(
                                        f"x.get({repr(safe_name)})"
                                    )
                            else:
                                match type_name:
                                    case ListTypeExpr(inner_type_name):
                                        encoder_name = TypeName(
                                            f"encode_{render_literal_type(inner_type_name)}"
                                        )
                                        encoder_names.add(encoder_name)
                                        typeddict_encoder.append(
                                            dedent(
                                                f"""\
                                            [
                                                {render_literal_type(encoder_name)}(y)
                                                for y in x[{repr(name)}]
                                            ]
                                            """.rstrip()
                                            )
                                        )
                        else:
                            if name in prop.required:
                                typeddict_encoder.append(f"x[{repr(safe_name)}]")
                            else:
                                typeddict_encoder.append(f"x.get({repr(safe_name)})")

                if name == "$kind":
                    # If the field is a literal, the Python type-checker will complain
                    # about the constructor not being able to specify a value for it:
                    # You can't put `$kind="ok"` in the ctor because `$` is not a valid
                    # character in an identifier, and putting `**{"$kind":"ok"}` makes
                    # it not recognize the `"ok"` as being `Literal["ok"]`, so we're
                    # stuck with an impossible-to-construct object.
                    field_value = "..."
                    match type_name:
                        case LiteralTypeExpr(literal_value):
                            field_value = repr(literal_value)
                    if name not in type.required:
                        value = ""
                        if base_model != "TypedDict":
                            value = dedent(
                                f"""\
                                = Field(
                                  default=None,
                                  alias={repr(name)}, # type: ignore
                                )
                                """
                            )
                        current_chunks.append(
                            f"  kind: {
                                render_type_expr(
                                    UnionTypeExpr(
                                        [
                                            type_name,
                                            NoneTypeExpr(),
                                        ]
                                    )
                                )
                            }{value}"
                        )
                    else:
                        value = ""
                        if base_model != "TypedDict":
                            value = dedent(
                                f"""\
                                = Field(
                                    {field_value},
                                    alias={repr(name)}, # type: ignore
                                )
                                """
                            )
                        current_chunks.append(
                            f"  kind: {render_type_expr(type_name)}{value}"
                        )
                else:
                    if name not in type.required:
                        if base_model == "TypedDict":
                            current_chunks.append(
                                reindent(
                                    "  ",
                                    f"""\
                        {name}: NotRequired[{
                                        render_type_expr(
                                            UnionTypeExpr([type_name, NoneTypeExpr()])
                                        )
                                    }]
                                """,
                                )
                            )
                        else:
                            current_chunks.append(
                                reindent(
                                    "  ",
                                    f"""\
                        {name}: {
                                        render_type_expr(
                                            UnionTypeExpr(
                                                [
                                                    type_name,
                                                    NoneTypeExpr(),
                                                ]
                                            )
                                        )
                                    } = None
                                """,
                                )
                            )
                    else:
                        current_chunks.append(
                            f"  {name}: {render_type_expr(type_name)}"
                        )
                typeddict_encoder.append(",")
            typeddict_encoder.append("}")
            # exclude_none
            typeddict_encoder = (
                ["{k: v for (k, v) in ("]
                + typeddict_encoder
                + [").items() if v is not None}"]
            )
        else:
            typeddict_encoder.append("{}")
            current_chunks.append("  pass")
        current_chunks.append("")

        if base_model == "TypedDict":
            binding = "x" if needs_binding else "_"
            encoder_name = TypeName(f"encode_{render_literal_type(prefix)}")
            encoder_names.add(encoder_name)
            _field_name = render_literal_type(encoder_name)
            current_chunks.insert(
                0,
                FileContents(
                    "\n".join(
                        [
                            dedent(
                                f"""\
                            def {_field_name}(
                                {binding}: {repr(render_literal_type(prefix))},
                            ) -> Any:
                                return (
                                {
                                    reindent(
                                        "                                ",
                                        "\n".join(typeddict_encoder),
                                    )
                                }
                            )
                            """
                            )
                        ]
                    )
                ),
            )
        chunks.append(FileContents("\n".join(current_chunks)))

    return (prefix, in_module, chunks, encoder_names)


def generate_common_client(
    client_name: str,
    handshake_type: HandshakeType,
    handshake_chunks: Sequence[str],
    modules: list[tuple[ModuleName, ClassName]],
) -> FileContents:
    chunks: list[str] = [ROOT_FILE_HEADER]
    chunks.extend(
        [
            f"from .{model_name} import {class_name}"
            for model_name, class_name in sorted(modules, key=lambda kv: kv[1])
        ]
    )
    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[ModuleName, ClassName, dict[RenderedPath, FileContents]]:
    serdes: list[tuple[list[TypeName], list[ModuleName], list[FileContents]]] = []

    def _type_adapter_definition(
        type_adapter_name: TypeName,
        _type: TypeExpression,
        module_info: list[ModuleName],
    ) -> tuple[list[TypeName], list[ModuleName], list[FileContents]]:
        rendered_type_expr = render_type_expr(_type)
        return (
            [type_adapter_name],
            module_info,
            [
                FileContents(
                    dedent(f"""
                    {render_type_expr(type_adapter_name)}: TypeAdapter[Any] = (
                        TypeAdapter({rendered_type_expr})
                    )
                """)
                )
            ],
        )

    class_name = ClassName(f"{schema_name.title()}Service")
    current_chunks: list[str] = [
        dedent(
            f"""\
              class {class_name}:
                def __init__(self, client: river.Client[Any]):
                  self.client = client
            """
        ),
    ]
    for name, procedure in schema.procedures.items():
        module_names = [ModuleName(name)]
        init_type: TypeExpression | None = None
        if procedure.init:
            init_type, module_info, init_chunks, encoder_names = encode_type(
                procedure.init,
                TypeName(f"{name.title()}Init"),
                input_base_class,
                module_names,
                permit_unknown_members=False,
            )
            serdes.append(
                (
                    [extract_inner_type(init_type), *encoder_names],
                    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,
            )
        )
        serdes.append(
            _type_adapter_definition(
                input_type_type_adapter_name, input_type, module_info
            )
        )
        output_type, module_info, output_chunks, encoder_names = encode_type(
            procedure.output,
            TypeName(f"{name.title()}Output"),
            "BaseModel",
            module_names,
            permit_unknown_members=True,
        )
        output_type_name = extract_inner_type(output_type)
        serdes.append(
            (
                [output_type_name, *encoder_names],
                module_info,
                output_chunks,
            )
        )
        output_type_type_adapter_name = TypeName(
            f"{render_literal_type(output_type_name)}TypeAdapter"
        )
        serdes.append(
            _type_adapter_definition(
                output_type_type_adapter_name, output_type, module_info
            )
        )
        output_module_info = module_info
        if procedure.errors:
            error_type, module_info, errors_chunks, encoder_names = encode_type(
                procedure.errors,
                TypeName(f"{name.title()}Errors"),
                "RiverError",
                module_names,
                permit_unknown_members=True,
            )
            if isinstance(error_type, NoneTypeExpr):
                error_type_name = TypeName("RiverError")
                error_type = error_type_name
            else:
                error_type_name = extract_inner_type(error_type)
                serdes.append(([error_type_name], module_info, errors_chunks))

        else:
            error_type_name = TypeName("RiverError")
            error_type = error_type_name

        error_type_type_adapter_name = TypeName(
            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
                )
            )
        output_or_error_type = UnionTypeExpr([output_type, error_type_name])

        # NB: These strings must be indented to at least the same level of
        #     the function strings in the branches below, otherwise `dedent`
        #     will pick our indentation level for normalization, which will
        #     break the "def" indentation presuppositions.
        output_type_adapter = render_literal_type(output_type_type_adapter_name)
        parse_output_method = f"""\
                            lambda x: {output_type_adapter}
                                .validate_python(
                                    x # type: ignore[arg-type]
                                )
                            """
        error_type_adapter = render_literal_type(error_type_type_adapter_name)
        parse_error_method = f"""\
                            lambda x: {error_type_adapter}
                                .validate_python(
                                    x # type: ignore[arg-type]
                                )
                            """

        # Init renderer
        render_init_method: str | None = None
        if init_type and procedure.init is not None:
            if input_base_class == "TypedDict":
                if is_literal(procedure.init):
                    render_init_method = "lambda x: x"
                elif isinstance(
                    procedure.init, RiverConcreteType
                ) and procedure.init.type in ["array"]:
                    match init_type:
                        case ListTypeExpr(init_type_name):
                            render_init_method = (
                                f"lambda xs: [encode_{init_type_name}(x) for x in xs]"
                            )
                else:
                    render_init_method = f"encode_{render_literal_type(init_type)}"
            else:
                init_type_name = extract_inner_type(init_type)
                init_type_type_adapter_name = TypeName(
                    f"{init_type_name.value}TypeAdapter"
                )
                serdes.append(
                    _type_adapter_definition(
                        init_type_type_adapter_name, init_type, 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: {input_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,
                              )
                            """

        assert render_input_method, (
            f"Unable to derive the input encoder from: {input_type}"
        )

        if isinstance(output_type, NoneTypeExpr):
            parse_output_method = "lambda x: None"

        if procedure.type == "rpc":
            current_chunks.extend(
                [
                    reindent(
                        "  ",
                        f"""\
            async def {name}(
              self,
              input: {render_type_expr(input_type)},
              timeout: datetime.timedelta,
            ) -> {render_type_expr(output_type)}:
              return await self.client.send_rpc(
                {repr(schema_name)},
                {repr(name)},
                input,
                {reindent("                    ", render_input_method)},
                {reindent("                    ", parse_output_method)},
                {reindent("                    ", parse_error_method)},
                timeout,
              )
                    """,
                    )
                ]
            )
        elif procedure.type == "subscription":
            output_or_error_type = UnionTypeExpr(
                [
                    output_or_error_type,
                    TypeName("RiverError"),
                ]
            )
            current_chunks.extend(
                [
                    reindent(
                        "  ",
                        f"""\
            async def {name}(
              self,
              input: {render_type_expr(input_type)},
            ) -> AsyncIterator[{render_type_expr(output_or_error_type)}]:
              return self.client.send_subscription(
                {repr(schema_name)},
                {repr(name)},
                input,
                {reindent("                    ", render_input_method)},
                {reindent("                    ", parse_output_method)},
                {reindent("                    ", parse_error_method)},
              )
                  """,
                    )
                ]
            )
        elif procedure.type == "upload":
            if init_type:
                assert render_init_method, "Expected an init renderer!"
                current_chunks.extend(
                    [
                        reindent(
                            "  ",
                            f"""\
            async def {name}(
              self,
              init: {render_type_expr(init_type)},
              inputStream: AsyncIterable[{render_type_expr(input_type)}],
            ) -> {render_type_expr(output_type)}:
              return await self.client.send_upload(
                {repr(schema_name)},
                {repr(name)},
                init,
                inputStream,
                {reindent("                    ", render_init_method)},
                {reindent("                    ", render_input_method)},
                {reindent("                    ", parse_output_method)},
                {reindent("                    ", parse_error_method)},
              )
                        """,
                        )
                    ]
                )
            else:
                current_chunks.extend(
                    [
                        reindent(
                            "  ",
                            f"""\
            async def {name}(
              self,
              inputStream: AsyncIterable[{render_type_expr(input_type)}],
            ) -> {render_type_expr(output_or_error_type)}:
              return await self.client.send_upload(
                {repr(schema_name)},
                {repr(name)},
                None,
                inputStream,
                None,
                {reindent("                    ", render_input_method)},
                {reindent("                    ", parse_output_method)},
                {reindent("                    ", parse_error_method)},
              )
                        """,
                        )
                    ]
                )
        elif procedure.type == "stream":
            output_or_error_type = UnionTypeExpr(
                [
                    output_or_error_type,
                    TypeName("RiverError"),
                ]
            )
            if init_type:
                assert render_init_method, "Expected an init renderer!"
                current_chunks.extend(
                    [
                        reindent(
                            "  ",
                            f"""\
            async def {name}(
              self,
              init: {render_type_expr(init_type)},
              inputStream: AsyncIterable[{render_type_expr(input_type)}],
            ) -> AsyncIterator[{render_type_expr(output_or_error_type)}]:
              return self.client.send_stream(
                {repr(schema_name)},
                {repr(name)},
                init,
                inputStream,
                {reindent("                    ", render_init_method)},
                {reindent("                    ", render_input_method)},
                {reindent("                    ", parse_output_method)},
                {reindent("                    ", parse_error_method)},
              )
                        """,
                        )
                    ]
                )
            else:
                current_chunks.extend(
                    [
                        reindent(
                            "  ",
                            f"""\
            async def {name}(
              self,
              inputStream: AsyncIterable[{render_type_expr(input_type)}],
            ) -> AsyncIterator[{render_type_expr(output_or_error_type)}]:
              return self.client.send_stream(
                {repr(schema_name)},
                {repr(name)},
                None,
                inputStream,
                None,
                {reindent("                    ", render_input_method)},
                {reindent("                    ", parse_output_method)},
                {reindent("                    ", parse_error_method)},
              )
                        """,
                        )
                    ]
                )

        current_chunks.append("")

    emitted_files: dict[RenderedPath, FileContents] = {}

    imports: dict[str, set[TypeName]] = dict()

    in_root: list[FileContents] = []
    for names, module_info, contents in serdes:
        if not module_info:
            in_root.extend(contents)
            continue

        imports.setdefault(".".join(module_info), set()).update(names)

        leaf_module = module_info[-1] + ".py"
        file_path = RenderedPath(str(Path(schema_name, *module_info[:-1], leaf_module)))
        existing = emitted_files.get(file_path, FileContents(FILE_HEADER))
        emitted_files[file_path] = FileContents("\n".join([existing] + contents))

    def render_names(xs: set[TypeName]) -> str:
        return ", ".join(sorted(render_literal_type(x) for x in xs))

    rendered_imports = [
        f"from .{dotted_modules} import {render_names(names)}"
        for dotted_modules, names in imports.items()
    ]

    emitted_files[RenderedPath(str(Path(f"{schema_name}/__init__.py")))] = FileContents(
        "\n".join([SERVICE_FILE_HEADER] + rendered_imports + in_root + current_chunks)
    )
    return (
        ModuleName(schema_name),
        class_name,
        emitted_files,
    )


def generate_river_client_module(
    client_name: str,
    schema_root: RiverSchema,
    typed_dict_inputs: bool,
) -> dict[RenderedPath, FileContents]:
    files: dict[RenderedPath, FileContents] = {}

    # Negotiate handshake shape
    handshake_chunks: list[str] = []
    if schema_root.handshakeSchema is not None:
        _handshake_type, _, contents, _ = encode_type(
            schema_root.handshakeSchema,
            TypeName("HandshakeSchema"),
            "BaseModel",
            [],
            permit_unknown_members=False,
        )
        handshake_chunks.extend(contents)
        handshake_type = HandshakeType(render_type_expr(_handshake_type))
    else:
        handshake_type = HandshakeType("Literal[None]")

    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():
        module_name, class_name, emitted_files = generate_individual_service(
            schema_name, schema, input_base_class
        )
        files.update(emitted_files)
        modules.append((module_name, class_name))

    main_contents = generate_common_client(
        client_name, handshake_type, handshake_chunks, modules
    )
    files[RenderedPath(str(Path("__init__.py")))] = main_contents

    return files


def schema_to_river_client_codegen(
    read_schema: Callable[[], TextIO],
    target_path: str,
    client_name: str,
    typed_dict_inputs: bool,
    file_opener: Callable[[Path], TextIO],
) -> 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
    ).items():
        module_path = Path(target_path).joinpath(subpath)
        module_path.parent.mkdir(mode=0o755, parents=True, exist_ok=True)
        with file_opener(module_path) as f:
            try:
                popen = subprocess.Popen(
                    ["ruff", "format", "-"],
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                )
                stdout, _ = popen.communicate(contents.encode())
                if popen.returncode != 0:
                    f.write(contents)
                f.write(stdout.decode("utf-8"))
            except:
                f.write(contents)
                raise