Specify type parameters for context in async grpc methods

Author: WouldYouKindly

mypy_protobuf/main.py

@@ -837,10 +837,10 @@ def write_grpc_iterator_type(self) -> None:
             )
         wl("")
 
-    def get_servicer_context_type(self) -> str:
+    def get_servicer_context_type(self, input_: str, output: str) -> str:
         """Get the type to use for the context parameter in servicer methods."""
         if self.grpc_type == GRPCType.ASYNC:
-            return self._import("grpc.aio", "ServicerContext")
+            return self._import("grpc.aio", f"ServicerContext[{input_}, {output}]")
         elif self.grpc_type == GRPCType.SYNC:
             return self._import("grpc", "ServicerContext")
         else:
@@ -928,19 +928,20 @@ def write_grpc_methods(self, service: d.ServiceDescriptorProto, scl_prefix: Sour
             wl("")
         for i, method in methods:
             scl = scl_prefix + [d.ServiceDescriptorProto.METHOD_FIELD_NUMBER, i]
+            input_type = self._servicer_input_type(method)
+            output_type = self._servicer_output_type(method)
 
             if self.generate_concrete_servicer_stubs is False:
                 wl("@{}", self._import("abc", "abstractmethod"))
             wl("def {}(", method.name)
             with self._indent():
                 wl("self,")
                 input_name = "request_iterator" if method.client_streaming else "request"
-                input_type = self._servicer_input_type(method)
                 wl(f"{input_name}: {input_type},")
-                wl("context: {},", self.get_servicer_context_type())
+                wl("context: {},", self.get_servicer_context_type(input_type, output_type))
             wl(
                 ") -> {}:{}",
-                self._servicer_output_type(method),
+                output_type,
                 " ..." if not self._has_comments(scl) else "",
             )
             if self._has_comments(scl):

test/generated-async-only/testproto/grpc/dummy_pb2_grpc.pyi

@@ -61,31 +61,31 @@ class DummyServiceServicer(metaclass=abc.ABCMeta):
     def UnaryUnary(
         self,
         request: testproto.grpc.dummy_pb2.DummyRequest,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[testproto.grpc.dummy_pb2.DummyRequest, collections.abc.Awaitable[testproto.grpc.dummy_pb2.DummyReply]],
     ) -> collections.abc.Awaitable[testproto.grpc.dummy_pb2.DummyReply]:
         """UnaryUnary"""
 
     @abc.abstractmethod
     def UnaryStream(
         self,
         request: testproto.grpc.dummy_pb2.DummyRequest,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[testproto.grpc.dummy_pb2.DummyRequest, collections.abc.AsyncIterator[testproto.grpc.dummy_pb2.DummyReply]],
     ) -> collections.abc.AsyncIterator[testproto.grpc.dummy_pb2.DummyReply]:
         """UnaryStream"""
 
     @abc.abstractmethod
     def StreamUnary(
         self,
         request_iterator: collections.abc.AsyncIterator[testproto.grpc.dummy_pb2.DummyRequest],
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[collections.abc.AsyncIterator[testproto.grpc.dummy_pb2.DummyRequest], collections.abc.Awaitable[testproto.grpc.dummy_pb2.DummyReply]],
     ) -> collections.abc.Awaitable[testproto.grpc.dummy_pb2.DummyReply]:
         """StreamUnary"""
 
     @abc.abstractmethod
     def StreamStream(
         self,
         request_iterator: collections.abc.AsyncIterator[testproto.grpc.dummy_pb2.DummyRequest],
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[collections.abc.AsyncIterator[testproto.grpc.dummy_pb2.DummyRequest], collections.abc.AsyncIterator[testproto.grpc.dummy_pb2.DummyReply]],
     ) -> collections.abc.AsyncIterator[testproto.grpc.dummy_pb2.DummyReply]:
         """StreamStream"""
 
@@ -117,15 +117,15 @@ class DeprecatedServiceServicer(metaclass=abc.ABCMeta):
     def DeprecatedMethod(
         self,
         request: testproto.grpc.dummy_pb2.DeprecatedRequest,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[testproto.grpc.dummy_pb2.DeprecatedRequest, collections.abc.Awaitable[testproto.grpc.dummy_pb2.DummyReply]],
     ) -> collections.abc.Awaitable[testproto.grpc.dummy_pb2.DummyReply]:
         """DeprecatedMethod"""
 
     @abc.abstractmethod
     def DeprecatedMethodNotDeprecatedRequest(
         self,
         request: testproto.grpc.dummy_pb2.DummyRequest,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[testproto.grpc.dummy_pb2.DummyRequest, collections.abc.Awaitable[testproto.grpc.dummy_pb2.DummyReply]],
     ) -> collections.abc.Awaitable[testproto.grpc.dummy_pb2.DummyReply]:
         """DeprecatedMethodNotDeprecatedRequest"""
 

test/generated-async-only/testproto/grpc/import_pb2_grpc.pyi

@@ -49,23 +49,23 @@ class SimpleServiceServicer(metaclass=abc.ABCMeta):
     def UnaryUnary(
         self,
         request: google.protobuf.empty_pb2.Empty,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[google.protobuf.empty_pb2.Empty, collections.abc.Awaitable[testproto.test_pb2.Simple1]],
     ) -> collections.abc.Awaitable[testproto.test_pb2.Simple1]:
         """UnaryUnary"""
 
     @abc.abstractmethod
     def UnaryStream(
         self,
         request: testproto.test_pb2.Simple1,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[testproto.test_pb2.Simple1, collections.abc.Awaitable[google.protobuf.empty_pb2.Empty]],
     ) -> collections.abc.Awaitable[google.protobuf.empty_pb2.Empty]:
         """UnaryStream"""
 
     @abc.abstractmethod
     def NoComment(
         self,
         request: testproto.test_pb2.Simple1,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[testproto.test_pb2.Simple1, collections.abc.Awaitable[google.protobuf.empty_pb2.Empty]],
     ) -> collections.abc.Awaitable[google.protobuf.empty_pb2.Empty]: ...
 
 def add_SimpleServiceServicer_to_server(servicer: SimpleServiceServicer, server: grpc.aio.Server) -> None: ...