From 3d0c065a28f5d2fa7a21c82dcb47d9c28c177320 Mon Sep 17 00:00:00 2001
From: Leandro Lucarella <luca-frequenz@llucax.com>
Date: Thu, 5 Sep 2024 11:34:30 +0200
Subject: [PATCH] Implement a new parameter to generate only async stubs

This commit adds a new parameter `async_only` to the plugin, that when
present, will only generate the async stubs.

The current {}AsyncStub classes are "fake", they are not present in the
.py file, so they can only be used to write type hints. Using them also
only works if the user manually casts the sync stub to the async stub.

Since most users will only want to use either the sync or the async
stub, this new parameter will allow to generate only the stubs that are
suited for the user's use case, and allow them to have less hacky code.

Signed-off-by: Leandro Lucarella <luca-frequenz@llucax.com>
---
 mypy_protobuf/main.py | 35 +++++++++++++++++++++--------------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/mypy_protobuf/main.py b/mypy_protobuf/main.py
index 47fe3e7e..78a3cde3 100644
--- a/mypy_protobuf/main.py
+++ b/mypy_protobuf/main.py
@@ -781,6 +781,7 @@ def write_grpc_services(
         self,
         services: Iterable[d.ServiceDescriptorProto],
         scl_prefix: SourceCodeLocation,
+        async_only: bool,
     ) -> None:
         wl = self._write_line
         for i, service in enumerate(services):
@@ -797,21 +798,25 @@ def write_grpc_services(
             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')}]"
+                # To support casting into FooAsyncStub, allow both Channel and aio.Channel here,
+                # but only if we are generating both sync and async stubs.
+                channel = self._import("grpc.aio", "Channel")
+                if not async_only:
+                    channel = f"{self._import('typing', 'Union')}[{self._import('grpc', 'Channel')}, {channel}]"
                 wl("def __init__(self, channel: {}) -> None: ...", channel)
-                self.write_grpc_stub_methods(service, scl)
+                self.write_grpc_stub_methods(service, scl, is_async=async_only)
 
-            # 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)
+            if not async_only:
+                # 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)
 
             # The service definition interface
             wl(
@@ -1009,6 +1014,7 @@ def generate_mypy_stubs(
 def generate_mypy_grpc_stubs(
     descriptors: Descriptors,
     response: plugin_pb2.CodeGeneratorResponse,
+    async_only: bool,
     quiet: bool,
     readable_stubs: bool,
     relax_strict_optional_primitives: bool,
@@ -1022,7 +1028,7 @@ def generate_mypy_grpc_stubs(
             grpc=True,
         )
         pkg_writer.write_grpc_async_hacks()
-        pkg_writer.write_grpc_services(fd.service, [d.FileDescriptorProto.SERVICE_FIELD_NUMBER])
+        pkg_writer.write_grpc_services(fd.service, [d.FileDescriptorProto.SERVICE_FIELD_NUMBER], async_only)
 
         assert name == fd.name
         assert fd.name.endswith(".proto")
@@ -1079,6 +1085,7 @@ def grpc() -> None:
         generate_mypy_grpc_stubs(
             Descriptors(request),
             response,
+            "async_only" in request.parameter,
             "quiet" in request.parameter,
             "readable_stubs" in request.parameter,
             "relax_strict_optional_primitives" in request.parameter,