Moving type ascriptions down to the method level

Author: blast-hardcheese

src/replit_river/v2/client.py

@@ -14,12 +14,6 @@
 
 from replit_river.v2.client_transport import ClientTransport
 from replit_river.error_schema import ERROR_CODE_UNKNOWN, RiverError, RiverException
-from replit_river.rpc import (
-    ErrorType,
-    InitType,
-    RequestType,
-    ResponseType,
-)
 from replit_river.transport_options import (
     HandshakeMetadataType,
     TransportOptions,
@@ -91,16 +85,16 @@ async def close(self) -> None:
     async def ensure_connected(self) -> None:
         await self._transport.get_or_create_session()
 
-    async def send_rpc(
+    async def send_rpc[R, A](
         self,
         service_name: str,
         procedure_name: str,
-        request: RequestType,
-        request_serializer: Callable[[RequestType], Any],
-        response_deserializer: Callable[[Any], ResponseType],
-        error_deserializer: Callable[[Any], ErrorType],
+        request: R,
+        request_serializer: Callable[[R], Any],
+        response_deserializer: Callable[[Any], A],
+        error_deserializer: Callable[[Any], RiverError],
         timeout: timedelta,
-    ) -> ResponseType:
+    ) -> A:
         with _trace_procedure("rpc", service_name, procedure_name) as span_handle:
             session = await self._transport.get_or_create_session()
             return await session.send_rpc(
@@ -114,17 +108,17 @@ async def send_rpc(
                 timeout,
             )
 
-    async def send_upload(
+    async def send_upload[I, R, A](
         self,
         service_name: str,
         procedure_name: str,
-        init: InitType | None,
-        request: AsyncIterable[RequestType],
-        init_serializer: Callable[[InitType], Any] | None,
-        request_serializer: Callable[[RequestType], Any],
-        response_deserializer: Callable[[Any], ResponseType],
-        error_deserializer: Callable[[Any], ErrorType],
-    ) -> ResponseType:
+        init: I | None,
+        request: AsyncIterable[R],
+        init_serializer: Callable[[I], Any] | None,
+        request_serializer: Callable[[R], Any],
+        response_deserializer: Callable[[Any], A],
+        error_deserializer: Callable[[Any], RiverError],
+    ) -> A:
         with _trace_procedure("upload", service_name, procedure_name) as span_handle:
             session = await self._transport.get_or_create_session()
             return await session.send_upload(
@@ -139,15 +133,15 @@ async def send_upload(
                 span_handle.span,
             )
 
-    async def send_subscription(
+    async def send_subscription[R, E, A](
         self,
         service_name: str,
         procedure_name: str,
-        request: RequestType,
-        request_serializer: Callable[[RequestType], Any],
-        response_deserializer: Callable[[Any], ResponseType],
-        error_deserializer: Callable[[Any], ErrorType],
-    ) -> AsyncGenerator[ResponseType | RiverError, None]:
+        request: R,
+        request_serializer: Callable[[R], Any],
+        response_deserializer: Callable[[Any], A],
+        error_deserializer: Callable[[Any], E],
+    ) -> AsyncGenerator[A | E, None]:
         with _trace_procedure(
             "subscription", service_name, procedure_name
         ) as span_handle:
@@ -165,17 +159,17 @@ async def send_subscription(
                     _record_river_error(span_handle, msg)
                 yield msg  # type: ignore # https://github.com/python/mypy/issues/10817
 
-    async def send_stream(
+    async def send_stream[I, R, E, A](
         self,
         service_name: str,
         procedure_name: str,
-        init: InitType | None,
-        request: AsyncIterable[RequestType],
-        init_serializer: Callable[[InitType], Any] | None,
-        request_serializer: Callable[[RequestType], Any],
-        response_deserializer: Callable[[Any], ResponseType],
-        error_deserializer: Callable[[Any], ErrorType],
-    ) -> AsyncGenerator[ResponseType | RiverError, None]:
+        init: I | None,
+        request: AsyncIterable[R],
+        init_serializer: Callable[[I], Any] | None,
+        request_serializer: Callable[[R], Any],
+        response_deserializer: Callable[[Any], A],
+        error_deserializer: Callable[[Any], E],
+    ) -> AsyncGenerator[A | E, None]:
         with _trace_procedure("stream", service_name, procedure_name) as span_handle:
             session = await self._transport.get_or_create_session()
             async for msg in session.send_stream(

src/replit_river/v2/client_session.py

@@ -15,6 +15,7 @@
 from replit_river.error_schema import (
     ERROR_CODE_CANCEL,
     ERROR_CODE_STREAM_CLOSED,
+    RiverError,
     RiverException,
     RiverServiceException,
     StreamClosedRiverServiceException,
@@ -28,10 +29,6 @@
     ACK_BIT,
     STREAM_CLOSED_BIT,
     STREAM_OPEN_BIT,
-    ErrorType,
-    InitType,
-    RequestType,
-    ResponseType,
 )
 from replit_river.seq_manager import (
     IgnoreMessageException,
@@ -168,17 +165,17 @@ async def _handle_messages_from_ws(self) -> None:
         except ConnectionClosed as e:
             raise e
 
-    async def send_rpc(
+    async def send_rpc[R, A](
         self,
         service_name: str,
         procedure_name: str,
-        request: RequestType,
-        request_serializer: Callable[[RequestType], Any],
-        response_deserializer: Callable[[Any], ResponseType],
-        error_deserializer: Callable[[Any], ErrorType],
+        request: R,
+        request_serializer: Callable[[R], Any],
+        response_deserializer: Callable[[Any], A],
+        error_deserializer: Callable[[Any], RiverError],
         span: Span,
         timeout: timedelta,
-    ) -> ResponseType:
+    ) -> A:
         """Sends a single RPC request to the server.
 
         Expects the input and output be messages that will be msgpacked.
@@ -233,18 +230,18 @@ async def send_rpc(
         except Exception as e:
             raise e
 
-    async def send_upload(
+    async def send_upload[I, R, A](
         self,
         service_name: str,
         procedure_name: str,
-        init: InitType | None,
-        request: AsyncIterable[RequestType],
-        init_serializer: Callable[[InitType], Any] | None,
-        request_serializer: Callable[[RequestType], Any],
-        response_deserializer: Callable[[Any], ResponseType],
-        error_deserializer: Callable[[Any], ErrorType],
+        init: I | None,
+        request: AsyncIterable[R],
+        init_serializer: Callable[[I], Any] | None,
+        request_serializer: Callable[[R], Any],
+        response_deserializer: Callable[[Any], A],
+        error_deserializer: Callable[[Any], RiverError],
         span: Span,
-    ) -> ResponseType:
+    ) -> A:
         """Sends an upload request to the server.
 
         Expects the input and output be messages that will be msgpacked.
@@ -320,16 +317,16 @@ async def send_upload(
         except Exception as e:
             raise e
 
-    async def send_subscription(
+    async def send_subscription[R, E, A](
         self,
         service_name: str,
         procedure_name: str,
-        request: RequestType,
-        request_serializer: Callable[[RequestType], Any],
-        response_deserializer: Callable[[Any], ResponseType],
-        error_deserializer: Callable[[Any], ErrorType],
+        request: R,
+        request_serializer: Callable[[R], Any],
+        response_deserializer: Callable[[Any], A],
+        error_deserializer: Callable[[Any], E],
         span: Span,
-    ) -> AsyncGenerator[ResponseType | ErrorType, None]:
+    ) -> AsyncGenerator[A | E, None]:
         """Sends a subscription request to the server.
 
         Expects the input and output be messages that will be msgpacked.
@@ -372,18 +369,18 @@ async def send_subscription(
         finally:
             output.close()
 
-    async def send_stream(
+    async def send_stream[I, R, E, A](
         self,
         service_name: str,
         procedure_name: str,
-        init: InitType | None,
-        request: AsyncIterable[RequestType],
-        init_serializer: Callable[[InitType], Any] | None,
-        request_serializer: Callable[[RequestType], Any],
-        response_deserializer: Callable[[Any], ResponseType],
-        error_deserializer: Callable[[Any], ErrorType],
+        init: I | None,
+        request: AsyncIterable[R],
+        init_serializer: Callable[[I], Any] | None,
+        request_serializer: Callable[[R], Any],
+        response_deserializer: Callable[[Any], A],
+        error_deserializer: Callable[[Any], E],
         span: Span,
-    ) -> AsyncGenerator[ResponseType | ErrorType, None]:
+    ) -> AsyncGenerator[A | E, None]:
         """Sends a subscription request to the server.
 
         Expects the input and output be messages that will be msgpacked.