Generate async-compatible stubs and servicers

RobinMcCorkell · RobinMcCorkell · commit 6a826d6ddd9b · 2023-02-09T13:58:00.000Z
diff --git a/mypy_protobuf/main.py b/mypy_protobuf/main.py
@@ -663,30 +663,77 @@ def _map_key_value_types(
 
         return ktype, vtype
 
-    def _callable_type(self, method: d.MethodDescriptorProto) -> str:
+    def _callable_type(self, method: d.MethodDescriptorProto, is_async: bool = False) -> str:
+        module = "grpc.aio" if is_async else "grpc"
         if method.client_streaming:
             if method.server_streaming:
-                return self._import("grpc", "StreamStreamMultiCallable")
+                return self._import(module, "StreamStreamMultiCallable")
             else:
-                return self._import("grpc", "StreamUnaryMultiCallable")
+                return self._import(module, "StreamUnaryMultiCallable")
         else:
             if method.server_streaming:
-                return self._import("grpc", "UnaryStreamMultiCallable")
+                return self._import(module, "UnaryStreamMultiCallable")
             else:
-                return self._import("grpc", "UnaryUnaryMultiCallable")
+                return self._import(module, "UnaryUnaryMultiCallable")
 
-    def _input_type(self, method: d.MethodDescriptorProto, use_stream_iterator: bool = True) -> str:
+    def _input_type(self, method: d.MethodDescriptorProto) -> str:
         result = self._import_message(method.input_type)
-        if use_stream_iterator and method.client_streaming:
-            result = f"{self._import('collections.abc', 'Iterator')}[{result}]"
         return result
 
-    def _output_type(self, method: d.MethodDescriptorProto, use_stream_iterator: bool = True) -> str:
+    def _servicer_input_type(self, method: d.MethodDescriptorProto) -> str:
+        result = self._import_message(method.input_type)
+        if method.client_streaming:
+            # See write_grpc_async_hacks().
+            result = f"_MaybeAsyncIterator[{result}]"
+        return result
+
+    def _output_type(self, method: d.MethodDescriptorProto) -> str:
         result = self._import_message(method.output_type)
-        if use_stream_iterator and method.server_streaming:
-            result = f"{self._import('collections.abc', 'Iterator')}[{result}]"
         return result
 
+    def _servicer_output_type(self, method: d.MethodDescriptorProto) -> str:
+        result = self._import_message(method.output_type)
+        if method.server_streaming:
+            # Union[Iterator[Resp], AsyncIterator[Resp]] is subtyped by Iterator[Resp] and AsyncIterator[Resp].
+            # So both can be used in the covariant function return position.
+            iterator = f"{self._import('typing', 'Iterator')}[{result}]"
+            aiterator = f"{self._import('typing', 'AsyncIterator')}[{result}]"
+            result = f"{self._import('typing', 'Union')}[{iterator}, {aiterator}]"
+        else:
+            # Union[Resp, Awaitable[Resp]] is subtyped by Resp and Awaitable[Resp].
+            # So both can be used in the covariant function return position.
+            # Awaitable[Resp] is equivalent to async def.
+            awaitable = f"{self._import('typing', 'Awaitable')}[{result}]"
+            result = f"{self._import('typing', 'Union')}[{result}, {awaitable}]"
+        return result
+
+    def write_grpc_async_hacks(self) -> None:
+        wl = self._write_line
+        # _MaybeAsyncIterator[Req] is supertyped by Iterator[Req] and AsyncIterator[Req].
+        # So both can be used in the contravariant function parameter position.
+        wl("_T = {}('_T')", self._import("typing", "TypeVar"))
+        wl("")
+        wl(
+            "class _MaybeAsyncIterator({}[_T], {}[_T], metaclass={}):",
+            self._import("typing", "AsyncIterator"),
+            self._import("typing", "Iterator"),
+            self._import("abc", "ABCMeta"),
+        )
+        with self._indent():
+            wl("...")
+        wl("")
+
+        # _ServicerContext is supertyped by grpc.ServicerContext and grpc.aio.ServicerContext
+        # So both can be used in the contravariant function parameter position.
+        wl(
+            "class _ServicerContext({}, {}):  # type: ignore",
+            self._import("grpc", "ServicerContext"),
+            self._import("grpc.aio", "ServicerContext"),
+        )
+        with self._indent():
+            wl("...")
+        wl("")
+
     def write_grpc_methods(self, service: d.ServiceDescriptorProto, scl_prefix: SourceCodeLocation) -> None:
         wl = self._write_line
         methods = [(i, m) for i, m in enumerate(service.method) if m.name not in PYTHON_RESERVED]
@@ -701,20 +748,20 @@ def write_grpc_methods(self, service: d.ServiceDescriptorProto, scl_prefix: Sour
             with self._indent():
                 wl("self,")
                 input_name = "request_iterator" if method.client_streaming else "request"
-                input_type = self._input_type(method)
+                input_type = self._servicer_input_type(method)
                 wl(f"{input_name}: {input_type},")
-                wl("context: {},", self._import("grpc", "ServicerContext"))
+                wl("context: _ServicerContext,")
             wl(
                 ") -> {}:{}",
-                self._output_type(method),
+                self._servicer_output_type(method),
                 " ..." if not self._has_comments(scl) else "",
             )
             if self._has_comments(scl):
                 with self._indent():
                     if not self._write_comments(scl):
                         wl("...")
 
-    def write_grpc_stub_methods(self, service: d.ServiceDescriptorProto, scl_prefix: SourceCodeLocation) -> None:
+    def write_grpc_stub_methods(self, service: d.ServiceDescriptorProto, scl_prefix: SourceCodeLocation, is_async: bool = False) -> None:
         wl = self._write_line
         methods = [(i, m) for i, m in enumerate(service.method) if m.name not in PYTHON_RESERVED]
         if not methods:
@@ -723,10 +770,10 @@ def write_grpc_stub_methods(self, service: d.ServiceDescriptorProto, scl_prefix:
         for i, method in methods:
             scl = scl_prefix + [d.ServiceDescriptorProto.METHOD_FIELD_NUMBER, i]
 
-            wl("{}: {}[", method.name, self._callable_type(method))
+            wl("{}: {}[", method.name, self._callable_type(method, is_async=is_async))
             with self._indent():
-                wl("{},", self._input_type(method, False))
-                wl("{},", self._output_type(method, False))
+                wl("{},", self._input_type(method))
+                wl("{},", self._output_type(method))
             wl("]")
             self._write_comments(scl)
 
@@ -743,17 +790,34 @@ def write_grpc_services(
             scl = scl_prefix + [i]
 
             # The stub client
-            wl(f"class {service.name}Stub:")
+            wl(
+                "class {}Stub:",
+                service.name,
+            )
             with self._indent():
                 if self._write_comments(scl):
                     wl("")
+                # To support casting into FooAsyncStub, allow both Channel and aio.Channel here.
+                channel = f"{self._import('typing', 'Union')}[{self._import('grpc', 'Channel')}, {self._import('grpc.aio', 'Channel')}]"
                 wl(
                     "def __init__(self, channel: {}) -> None: ...",
-                    self._import("grpc", "Channel"),
+                    channel
                 )
                 self.write_grpc_stub_methods(service, scl)
             wl("")
 
+            # The (fake) async stub client
+            wl(
+                "class {}AsyncStub:",
+                service.name,
+            )
+            with self._indent():
+                if self._write_comments(scl):
+                    wl("")
+                # No __init__ since this isn't a real class (yet), and requires manual casting to work.
+                self.write_grpc_stub_methods(service, scl, is_async=True)
+            wl("")
+
             # The service definition interface
             wl(
                 "class {}Servicer(metaclass={}):",
@@ -765,11 +829,13 @@ def write_grpc_services(
                     wl("")
                 self.write_grpc_methods(service, scl)
             wl("")
+            server = self._import('grpc', 'Server')
+            aserver = self._import('grpc.aio', 'Server')
             wl(
                 "def add_{}Servicer_to_server(servicer: {}Servicer, server: {}) -> None: ...",
                 service.name,
                 service.name,
-                self._import("grpc", "Server"),
+                f"{self._import('typing', 'Union')}[{server}, {aserver}]",
             )
             wl("")
 
@@ -960,6 +1026,7 @@ def generate_mypy_grpc_stubs(
             relax_strict_optional_primitives,
             grpc=True,
         )
+        pkg_writer.write_grpc_async_hacks()
         pkg_writer.write_grpc_services(fd.service, [d.FileDescriptorProto.SERVICE_FIELD_NUMBER])
 
         assert name == fd.name
diff --git a/stubtest_allowlist.txt b/stubtest_allowlist.txt
@@ -48,6 +48,10 @@ testproto.readme_enum_pb2._?MyEnum(EnumTypeWrapper)?
 testproto.nested.nested_pb2.AnotherNested._?NestedEnum(EnumTypeWrapper)?
 testproto.nested.nested_pb2.AnotherNested.NestedMessage._?NestedEnum2(EnumTypeWrapper)?
 
+# Our fake async stubs are not there at runtime (yet)
+testproto.grpc.dummy_pb2_grpc.DummyServiceAsyncStub
+testproto.grpc.import_pb2_grpc.SimpleServiceAsyncStub
+
 # Part of an "EXPERIMENTAL API" according to comment. Not documented.
 testproto.grpc.dummy_pb2_grpc.DummyService
 testproto.grpc.import_pb2_grpc.SimpleService
diff --git a/test/generated/testproto/grpc/dummy_pb2_grpc.pyi b/test/generated/testproto/grpc/dummy_pb2_grpc.pyi
@@ -3,14 +3,23 @@
 isort:skip_file
 https://github.com/vmagamedov/grpclib/blob/master/tests/dummy.proto"""
 import abc
-import collections.abc
 import grpc
+import grpc.aio
 import testproto.grpc.dummy_pb2
+import typing
+
+_T = typing.TypeVar('_T')
+
+class _MaybeAsyncIterator(typing.AsyncIterator[_T], typing.Iterator[_T], metaclass=abc.ABCMeta):
+    ...
+
+class _ServicerContext(grpc.ServicerContext, grpc.aio.ServicerContext):  # type: ignore
+    ...
 
 class DummyServiceStub:
     """DummyService"""
 
-    def __init__(self, channel: grpc.Channel) -> None: ...
+    def __init__(self, channel: typing.Union[grpc.Channel, grpc.aio.Channel]) -> None: ...
     UnaryUnary: grpc.UnaryUnaryMultiCallable[
         testproto.grpc.dummy_pb2.DummyRequest,
         testproto.grpc.dummy_pb2.DummyReply,
@@ -32,36 +41,60 @@ class DummyServiceStub:
     ]
     """StreamStream"""
 
+class DummyServiceAsyncStub:
+    """DummyService"""
+
+    UnaryUnary: grpc.aio.UnaryUnaryMultiCallable[
+        testproto.grpc.dummy_pb2.DummyRequest,
+        testproto.grpc.dummy_pb2.DummyReply,
+    ]
+    """UnaryUnary"""
+    UnaryStream: grpc.aio.UnaryStreamMultiCallable[
+        testproto.grpc.dummy_pb2.DummyRequest,
+        testproto.grpc.dummy_pb2.DummyReply,
+    ]
+    """UnaryStream"""
+    StreamUnary: grpc.aio.StreamUnaryMultiCallable[
+        testproto.grpc.dummy_pb2.DummyRequest,
+        testproto.grpc.dummy_pb2.DummyReply,
+    ]
+    """StreamUnary"""
+    StreamStream: grpc.aio.StreamStreamMultiCallable[
+        testproto.grpc.dummy_pb2.DummyRequest,
+        testproto.grpc.dummy_pb2.DummyReply,
+    ]
+    """StreamStream"""
+
 class DummyServiceServicer(metaclass=abc.ABCMeta):
     """DummyService"""
 
     @abc.abstractmethod
     def UnaryUnary(
         self,
         request: testproto.grpc.dummy_pb2.DummyRequest,
-        context: grpc.ServicerContext,
-    ) -> testproto.grpc.dummy_pb2.DummyReply:
+        context: _ServicerContext,
+    ) -> typing.Union[testproto.grpc.dummy_pb2.DummyReply, typing.Awaitable[testproto.grpc.dummy_pb2.DummyReply]]:
         """UnaryUnary"""
     @abc.abstractmethod
     def UnaryStream(
         self,
         request: testproto.grpc.dummy_pb2.DummyRequest,
-        context: grpc.ServicerContext,
-    ) -> collections.abc.Iterator[testproto.grpc.dummy_pb2.DummyReply]:
+        context: _ServicerContext,
+    ) -> typing.Union[typing.Iterator[testproto.grpc.dummy_pb2.DummyReply], typing.AsyncIterator[testproto.grpc.dummy_pb2.DummyReply]]:
         """UnaryStream"""
     @abc.abstractmethod
     def StreamUnary(
         self,
-        request_iterator: collections.abc.Iterator[testproto.grpc.dummy_pb2.DummyRequest],
-        context: grpc.ServicerContext,
-    ) -> testproto.grpc.dummy_pb2.DummyReply:
+        request_iterator: _MaybeAsyncIterator[testproto.grpc.dummy_pb2.DummyRequest],
+        context: _ServicerContext,
+    ) -> typing.Union[testproto.grpc.dummy_pb2.DummyReply, typing.Awaitable[testproto.grpc.dummy_pb2.DummyReply]]:
         """StreamUnary"""
     @abc.abstractmethod
     def StreamStream(
         self,
-        request_iterator: collections.abc.Iterator[testproto.grpc.dummy_pb2.DummyRequest],
-        context: grpc.ServicerContext,
-    ) -> collections.abc.Iterator[testproto.grpc.dummy_pb2.DummyReply]:
+        request_iterator: _MaybeAsyncIterator[testproto.grpc.dummy_pb2.DummyRequest],
+        context: _ServicerContext,
+    ) -> typing.Union[typing.Iterator[testproto.grpc.dummy_pb2.DummyReply], typing.AsyncIterator[testproto.grpc.dummy_pb2.DummyReply]]:
         """StreamStream"""
 
-def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: grpc.Server) -> None: ...
+def add_DummyServiceServicer_to_server(servicer: DummyServiceServicer, server: typing.Union[grpc.Server, grpc.aio.Server]) -> None: ...
diff --git a/test/generated/testproto/grpc/import_pb2_grpc.pyi b/test/generated/testproto/grpc/import_pb2_grpc.pyi
@@ -5,12 +5,22 @@ isort:skip_file
 import abc
 import google.protobuf.empty_pb2
 import grpc
+import grpc.aio
 import testproto.test_pb2
+import typing
+
+_T = typing.TypeVar('_T')
+
+class _MaybeAsyncIterator(typing.AsyncIterator[_T], typing.Iterator[_T], metaclass=abc.ABCMeta):
+    ...
+
+class _ServicerContext(grpc.ServicerContext, grpc.aio.ServicerContext):  # type: ignore
+    ...
 
 class SimpleServiceStub:
     """SimpleService"""
 
-    def __init__(self, channel: grpc.Channel) -> None: ...
+    def __init__(self, channel: typing.Union[grpc.Channel, grpc.aio.Channel]) -> None: ...
     UnaryUnary: grpc.UnaryUnaryMultiCallable[
         google.protobuf.empty_pb2.Empty,
         testproto.test_pb2.Simple1,
@@ -26,28 +36,46 @@ class SimpleServiceStub:
         google.protobuf.empty_pb2.Empty,
     ]
 
+class SimpleServiceAsyncStub:
+    """SimpleService"""
+
+    UnaryUnary: grpc.aio.UnaryUnaryMultiCallable[
+        google.protobuf.empty_pb2.Empty,
+        testproto.test_pb2.Simple1,
+    ]
+    """UnaryUnary"""
+    UnaryStream: grpc.aio.UnaryUnaryMultiCallable[
+        testproto.test_pb2.Simple1,
+        google.protobuf.empty_pb2.Empty,
+    ]
+    """UnaryStream"""
+    NoComment: grpc.aio.UnaryUnaryMultiCallable[
+        testproto.test_pb2.Simple1,
+        google.protobuf.empty_pb2.Empty,
+    ]
+
 class SimpleServiceServicer(metaclass=abc.ABCMeta):
     """SimpleService"""
 
     @abc.abstractmethod
     def UnaryUnary(
         self,
         request: google.protobuf.empty_pb2.Empty,
-        context: grpc.ServicerContext,
-    ) -> testproto.test_pb2.Simple1:
+        context: _ServicerContext,
+    ) -> typing.Union[testproto.test_pb2.Simple1, typing.Awaitable[testproto.test_pb2.Simple1]]:
         """UnaryUnary"""
     @abc.abstractmethod
     def UnaryStream(
         self,
         request: testproto.test_pb2.Simple1,
-        context: grpc.ServicerContext,
-    ) -> google.protobuf.empty_pb2.Empty:
+        context: _ServicerContext,
+    ) -> typing.Union[google.protobuf.empty_pb2.Empty, typing.Awaitable[google.protobuf.empty_pb2.Empty]]:
         """UnaryStream"""
     @abc.abstractmethod
     def NoComment(
         self,
         request: testproto.test_pb2.Simple1,
-        context: grpc.ServicerContext,
-    ) -> google.protobuf.empty_pb2.Empty: ...
+        context: _ServicerContext,
+    ) -> typing.Union[google.protobuf.empty_pb2.Empty, typing.Awaitable[google.protobuf.empty_pb2.Empty]]: ...
 
-def add_SimpleServiceServicer_to_server(servicer: SimpleServiceServicer, server: grpc.Server) -> None: ...
+def add_SimpleServiceServicer_to_server(servicer: SimpleServiceServicer, server: typing.Union[grpc.Server, grpc.aio.Server]) -> None: ...
diff --git a/test/test_grpc_async_usage.py b/test/test_grpc_async_usage.py
@@ -53,7 +53,7 @@ async def test_grpc() -> None:
     server = make_server()
     await server.start()
     async with grpc.aio.insecure_channel(ADDRESS) as channel:
-        client = dummy_pb2_grpc.DummyServiceStub(channel)
+        client: dummy_pb2_grpc.DummyServiceAsyncStub = dummy_pb2_grpc.DummyServiceStub(channel)  # type: ignore
         request = dummy_pb2.DummyRequest(value="cprg")
         result1 = await client.UnaryUnary(request)
         result2 = client.UnaryStream(dummy_pb2.DummyRequest(value=result1.value))
diff --git a/test_negative/output.expected.3.8 b/test_negative/output.expected.3.8
diff --git a/test_negative/output.expected.3.8.omit_linenos b/test_negative/output.expected.3.8.omit_linenos
