Improve typing

Makefile

@@ -0,0 +1,9 @@
+lint:
+	uv run black --check .
+	uv run ruff check .
+	uv run mypy .
+	uv run pyright-python .
+
+format:
+	uv run black .
+	uv run ruff check . --fix

pyproject.toml

@@ -47,6 +47,7 @@ dev-dependencies = [
     "ruff>=0.0.278",
     "types-protobuf>=4.24.0.20240311",
     "types-nanoid>=2.0.0.20240601",
+    "pyright>=1.1.389",
 ]
 
 [tool.ruff]

replit_river/rpc.py

@@ -127,7 +127,7 @@ def set(self, carrier: TransportMessage, key: str, value: str) -> None:
                 logger.warning("unknown trace propagation key", extra={"key": key})
 
 
-class GrpcContext(grpc.aio.ServicerContext):
+class GrpcContext(grpc.aio.ServicerContext, Generic[RequestType, ResponseType]):
     """Represents a gRPC-compatible ServicerContext for River interop."""
 
     def __init__(self, peer: str) -> None:
@@ -229,9 +229,8 @@ async def wrapped(
         input: Channel[Any],
         output: Channel[Any],
     ) -> None:
-        context = None
+        context: GrpcContext[RequestType, ResponseType] = GrpcContext(peer)
         try:
-            context = GrpcContext(peer)
             request = request_deserializer(await input.get())
             response = method(request, context)
             if isinstance(response, Awaitable):
@@ -287,9 +286,8 @@ async def wrapped(
         input: Channel[Any],
         output: Channel[Any],
     ) -> None:
-        context = None
+        context: GrpcContext[RequestType, ResponseType] = GrpcContext(peer)
         try:
-            context = GrpcContext(peer)
             request = request_deserializer(await input.get())
             iterator = method(request, context)
             if isinstance(iterator, AsyncIterable):
@@ -349,7 +347,7 @@ async def wrapped(
     ) -> None:
         task_manager = BackgroundTaskManager()
         try:
-            context = GrpcContext(peer)
+            context: GrpcContext[RequestType, ResponseType] = GrpcContext(peer)
             request: Channel[RequestType] = Channel(MAX_MESSAGE_BUFFER_SIZE)
 
             async def _convert_inputs() -> None:
@@ -426,9 +424,8 @@ async def wrapped(
         output: Channel[Any],
     ) -> None:
         task_manager = BackgroundTaskManager()
-        context = None
+        context: GrpcContext[RequestType, ResponseType] = GrpcContext(peer)
         try:
-            context = GrpcContext(peer)
             request: Channel[RequestType] = Channel(MAX_MESSAGE_BUFFER_SIZE)
 
             async def _convert_inputs() -> None:

tests/test_message_buffer.py

@@ -11,12 +11,12 @@ def mock_transport_message(seq: int) -> TransportMessage:
         seq=seq,
         id="test",
         ack=0,
-        from_="test",
+        from_="test",  # type: ignore
         to="test",
         streamId="test",
         controlFlags=0,
         payload=0,
-        model_config={},
+        model_config={},  # type: ignore
     )