more fixes

Author: jackyzha0

.github/workflows/ci.yaml

@@ -60,14 +60,24 @@ jobs:
 
       - name: Install Python dependencies
         working-directory: python-client
-        run: pip install -e ".[dev]"
+        run: pip install -e ".[dev]" ty
 
       - name: Python lint
         working-directory: python-client
         run: |
           ruff check .
           ruff format --check .
 
+      - name: Python type check
+        working-directory: python-client
+        run: ty check river/
+
       - name: Python tests
         working-directory: python-client
         run: python -m pytest tests/ -v
+
+      - name: Python type check generated clients
+        working-directory: python-client
+        run: |
+          ty check tests/generated/
+          ty check tests/test_codegen.py

python-client/pyproject.toml

@@ -29,10 +29,14 @@ testpaths = ["tests"]
 
 [tool.ruff]
 target-version = "py310"
+exclude = ["tests/generated"]
 
 [tool.ruff.lint]
 select = ["E", "F", "I", "W"]
 
+[tool.ty.environment]
+extra-paths = ["tests"]
+
 [tool.setuptools.packages.find]
 include = ["river*"]
 

python-client/river/__init__.py

@@ -1,13 +1,28 @@
-"""River protocol v2.0 Python client implementation."""
+"""River protocol v2.0 Python client implementation.
 
-from river.client import RiverClient
+This client was generated with the assistance of AI (Claude).
+"""
+
+from river.client import (
+    ErrResult,
+    OkResult,
+    RiverClient,
+    StreamResult,
+    SubscriptionResult,
+    UploadResult,
+)
 from river.codec import BinaryCodec, NaiveJsonCodec
 from river.streams import Readable, Writable
 from river.transport import WebSocketClientTransport
 from river.types import Err, Ok, TransportMessage
 
 __all__ = [
     "RiverClient",
+    "OkResult",
+    "ErrResult",
+    "StreamResult",
+    "UploadResult",
+    "SubscriptionResult",
     "WebSocketClientTransport",
     "NaiveJsonCodec",
     "BinaryCodec",

python-client/river/client.py

@@ -9,7 +9,9 @@
 import asyncio
 import logging
 from dataclasses import dataclass
-from typing import Any, Callable
+from typing import Any, Callable, Generic, Literal, TypeVar
+
+from typing_extensions import TypedDict
 
 from river.streams import Readable, Writable
 from river.transport import WebSocketClientTransport
@@ -29,28 +31,43 @@
 
 logger = logging.getLogger(__name__)
 
+T = TypeVar("T")
+TPayload = TypeVar("TPayload")
 
-@dataclass
-class RpcResult:
-    """Result of an RPC call."""
 
-    ok: bool
-    payload: Any
+class OkResult(TypedDict, Generic[TPayload]):
+    """Successful result from a procedure call."""
+
+    ok: Literal[True]
+    payload: TPayload
+
+
+class ErrResult(TypedDict, Generic[TPayload]):
+    """Error result from a procedure call."""
+
+    ok: Literal[False]
+    payload: TPayload
 
 
 @dataclass
-class StreamResult:
-    """Result of opening a stream procedure."""
+class StreamResult(Generic[T]):
+    """Result of opening a stream procedure.
 
-    req_writable: Writable
+    Generic over the input type ``T`` written to ``req_writable``.
+    """
+
+    req_writable: Writable[T]
     res_readable: Readable
 
 
 @dataclass
-class UploadResult:
-    """Result of opening an upload procedure."""
+class UploadResult(Generic[T]):
+    """Result of opening an upload procedure.
 
-    req_writable: Writable
+    Generic over the input type ``T`` written to ``req_writable``.
+    """
+
+    req_writable: Writable[T]
     finalize: Callable[[], Any]  # async callable returning RpcResult
 
 
@@ -113,7 +130,7 @@ async def rpc(
         procedure_name: str,
         init: Any,
         abort_signal: asyncio.Event | None = None,
-    ) -> dict[str, Any]:
+    ) -> Any:
         """Invoke an RPC procedure.
 
         Returns the result dict: {"ok": True/False, "payload": ...}

python-client/river/codec.py

@@ -78,18 +78,18 @@ class BinaryCodec(Codec):
     name = "binary"
 
     def to_buffer(self, obj: dict[str, Any]) -> bytes:
-        import msgpack  # type: ignore[import-untyped]
+        import msgpack
 
         return msgpack.packb(obj, use_bin_type=True, default=self._ext_encode)
 
     def from_buffer(self, buf: bytes) -> dict[str, Any]:
-        import msgpack  # type: ignore[import-untyped]
+        import msgpack
 
         return msgpack.unpackb(buf, raw=False, ext_hook=self._ext_decode)
 
     @staticmethod
     def _ext_encode(obj: Any) -> Any:
-        import msgpack  # type: ignore[import-untyped]
+        import msgpack
 
         if isinstance(obj, int) and (obj > _MSGPACK_INT_MAX or obj < _MSGPACK_INT_MIN):
             # Encode as string in extension type 0 (matches TS BigInt ext)
@@ -99,7 +99,7 @@ def _ext_encode(obj: Any) -> Any:
 
     @staticmethod
     def _ext_decode(code: int, data: bytes) -> Any:
-        import msgpack  # type: ignore[import-untyped]
+        import msgpack
 
         if code == _BIGINT_EXT_TYPE:
             val = msgpack.unpackb(data, raw=False)

python-client/river/codegen/emitter.py

@@ -29,6 +29,19 @@
 _env.filters["pascal"] = _to_pascal_case
 
 
+def _result_type(proc) -> str:  # noqa: ANN001
+    """Build the typed result annotation for a procedure."""
+    ok = f"OkResult[{proc.output_type.annotation}]"
+    if proc.error_type:
+        err = f"ErrResult[{proc.error_type.annotation} | ProtocolError]"
+    else:
+        err = "ErrResult[ProtocolError]"
+    return f"{ok} | {err}"
+
+
+_env.filters["result_type"] = _result_type
+
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -51,6 +64,9 @@ def _collect_used_type_names(svc: ServiceDef, ir: SchemaIR) -> list[str]:
         _extract_names(proc.init_type.annotation, td_names, names)
         if proc.input_type:
             _extract_names(proc.input_type.annotation, td_names, names)
+        _extract_names(proc.output_type.annotation, td_names, names)
+        if proc.error_type:
+            _extract_names(proc.error_type.annotation, td_names, names)
 
     return sorted(names)
 
@@ -110,22 +126,20 @@ def render_service_client(svc: ServiceDef, ir: SchemaIR, import_prefix: str) ->
     type_names = _collect_used_type_names(svc, ir)
     types_module = "._types" if import_prefix == "." else f"{import_prefix}_types"
 
-    needs_readable = any(
-        p.proc_type in ("stream", "subscription") for p in svc.procedures
-    )
-    needs_writable = any(p.proc_type in ("stream", "upload") for p in svc.procedures)
-
-    wrappers = [
-        p for p in svc.procedures if p.proc_type in ("stream", "upload", "subscription")
-    ]
+    proc_types = {p.proc_type for p in svc.procedures}
+    has_rpc = "rpc" in proc_types
+    has_stream = "stream" in proc_types
+    has_upload = "upload" in proc_types
+    has_subscription = "subscription" in proc_types
 
     return _env.get_template("service_client.py.j2").render(
         service=svc,
         type_names=type_names,
         types_module=types_module,
-        needs_readable=needs_readable,
-        needs_writable=needs_writable,
-        wrappers=wrappers,
+        has_rpc=has_rpc,
+        has_stream=has_stream,
+        has_upload=has_upload,
+        has_subscription=has_subscription,
     )
 
 

python-client/river/codegen/schema.py

@@ -118,10 +118,23 @@ def _to_snake_case(s: str) -> str:
 
 
 def _safe_field_name(name: str) -> str:
-    """Ensure a field name is a valid Python identifier."""
-    name = _sanitize_identifier(name)
+    """Ensure a field name is a valid Python identifier.
+
+    Raises ValueError if the name requires sanitization that would
+    change it from its wire representation, since TypedDict keys must
+    match the dict keys sent on the wire.
+    """
+    sanitized = _sanitize_identifier(name)
+    if sanitized != name:
+        raise ValueError(
+            f"schema property {name!r} is not a valid Python identifier "
+            f"and cannot be represented in a TypedDict"
+        )
     if keyword.iskeyword(name):
-        return name + "_"
+        raise ValueError(
+            f"schema property {name!r} is a Python keyword "
+            f"and cannot be used as a TypedDict field"
+        )
     return name