Overload __new__ (#710)

* Overload __new__

Mypy can't handle a large number of generic overloads (python/mypy#19622) so instead generate 2 different stub classes and override __new__. Seems to have the same functionality, without the generics overhead. Minor feature change in the revealed type of a stub class.

Signed-off-by: Aidan Jensen <aidandj.github@gmail.com>

* Write union methods to main typestub. Add

Signed-off-by: Aidan Jensen <aidandj.github@gmail.com>

* Clean up generated code and readme

* Document async usage
* Add tests for async/sync attribute usage
* Generate sync methods by default, and type ignore async override for new override

Signed-off-by: Aidan Jensen <aidandj.github@gmail.com>

* Fix missing method comment regression

Signed-off-by: Aidan Jensen <aidandj.github@gmail.com>

---------

Signed-off-by: Aidan Jensen <aidandj.github@gmail.com>

Author: aidandj

CHANGELOG.md

@@ -10,6 +10,8 @@
 - Switch to types-grpcio instead of no longer maintained grpc-stubs
 - Add `_HasFieldArgType` and `_ClearFieldArgType` aliases to allow for typing field manipulation functions
 - Add `_WhichOneofArgType_<oneof_name>` and `_WhichOneofReturnType_<oneof_name>` type aliases
+- Use `__new__` overloads for async stubs instead of `TypeVar` based `__init__` overloads.
+  - https://github.com/nipunn1313/mypy-protobuf/issues/707
 
 ## 3.7.0
 

README.md

@@ -371,6 +371,32 @@ protoc \
 Note that generated code for grpc will work only together with code for python and locations should be the same.
 If you need stubs for grpc internal code we suggest using this package https://pypi.org/project/types-grpcio/
 
+#### Async GRPC usage
+
+`mypy-protobuf` generates stubs that are compatible with both sync and async usage, with a few caveats.
+
+In a simple use case, the stubs work as expected.
+
+```python
+stub = dummy_Pb2_grpc.DummyServiceStub(grpc.aio.insecure_channel("localhost:1234"))
+result = await stub.UnaryUnary(dummy_pb2.DummyRequest(value="cprg"))
+typing.assert_type(result, dummy_pb2.DummyReply)
+```
+
+If you need to explicitly type something as an async stub (class attr, etc) then you must use deferred annotations, and the async stub, as it does not exist at runtime.
+
+```python
+class TestAttribute:
+    stub: "dummy_pb2_grpc.DummyServiceAsyncStub"
+
+    def __init__(self) -> None:
+        self.stub = dummy_pb2_grpc.DummyServiceStub(grpc.aio.insecure_channel("localhost:1234"))
+
+    async def test(self) -> None:
+        result = await self.stub.UnaryUnary(dummy_pb2.DummyRequest(value="cprg"))
+        typing.assert_type(result, dummy_pb2.DummyReply)
+```
+
 ### `_ClearFieldArgType`, `_WhichOneofArgType_<oneof_name>`, `_WhichOneofReturnType_<oneof_name>` and `_HasFieldArgType` aliases
 
 Where applicable, type aliases are generated for the arguments to `ClearField`, `WhichOneof` and `HasField`. These can be used to create typed functions for field manipulation:

mypy_protobuf/main.py

@@ -87,11 +87,6 @@
 }
 
 
-def _build_typevar_name(service_name: str, method_name: str) -> str:
-    # Prefix with underscore to avoid public api error: https://stackoverflow.com/a/78871465
-    return f"_{service_name}{method_name}Type"
-
-
 def _mangle_global_identifier(name: str) -> str:
     """
     Module level identifiers are mangled and aliased so that they can be disambiguated
@@ -180,7 +175,7 @@ def _import(self, path: str, name: str) -> str:
         eg. self._import("typing", "Literal") -> "Literal"
         """
         if path == "typing_extensions":
-            stabilization = {"TypeAlias": (3, 10), "TypeVar": (3, 13)}
+            stabilization = {"TypeAlias": (3, 10), "TypeVar": (3, 13), "type_check_only": (3, 12)}
             assert name in stabilization
             if not self.typing_extensions_min or self.typing_extensions_min < stabilization[name]:
                 self.typing_extensions_min = stabilization[name]
@@ -816,45 +811,28 @@ def write_grpc_async_hacks(self) -> None:
             wl("...")
         wl("")
 
-    def write_grpc_type_vars(self, service: d.ServiceDescriptorProto) -> None:
+    def write_grpc_stub_methods(self, service: d.ServiceDescriptorProto, scl_prefix: SourceCodeLocation, *, is_async: bool, both: bool = False, ignore_assignment_errors: 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:
             return
-        for _, method in methods:
-            wl("{} = {}(", _build_typevar_name(service.name, method.name), self._import("typing_extensions", "TypeVar"))
-            with self._indent():
-                wl("'{}',", _build_typevar_name(service.name, method.name))
-                wl("{}[", self._callable_type(method, is_async=False))
-                with self._indent():
-                    wl("{},", self._input_type(method))
-                    wl("{},", self._output_type(method))
-                wl("],")
-                wl("{}[", self._callable_type(method, is_async=True))
-                with self._indent():
-                    wl("{},", self._input_type(method))
-                    wl("{},", self._output_type(method))
-                wl("],")
-                wl("default={}[", self._callable_type(method, is_async=False))
-                with self._indent():
-                    wl("{},", self._input_type(method))
-                    wl("{},", self._output_type(method))
-                wl("],")
-            wl(")")
-            wl("")
 
-    def write_self_types(self, service: d.ServiceDescriptorProto, is_async: bool) -> 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:
-            return
-        for _, method in methods:
-            with self._indent():
-                wl("{}[", self._callable_type(method, is_async=is_async))
-                with self._indent():
-                    wl("{},", self._input_type(method))
-                    wl("{},", self._output_type(method))
-                wl("],")
+        def type_str(method: d.MethodDescriptorProto, is_async: bool) -> str:
+            return f"{self._callable_type(method, is_async=is_async)}[{self._input_type(method)}, {self._output_type(method)}]"
+
+        for i, method in methods:
+            scl = scl_prefix + [d.ServiceDescriptorProto.METHOD_FIELD_NUMBER, i]
+            if both:
+                wl(
+                    "{}: {}[{}, {}]",
+                    method.name,
+                    self._import("typing", "Union"),
+                    type_str(method, is_async=False),
+                    type_str(method, is_async=True),
+                )
+            else:
+                wl("{}: {}{}", method.name, type_str(method, is_async=is_async), "" if not ignore_assignment_errors else "  # type: ignore[assignment]")
+            self._write_comments(scl)
 
     def write_grpc_methods(self, service: d.ServiceDescriptorProto, scl_prefix: SourceCodeLocation) -> None:
         wl = self._write_line
@@ -885,19 +863,6 @@ def write_grpc_methods(self, service: d.ServiceDescriptorProto, scl_prefix: Sour
                         wl("...")
             wl("")
 
-    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:
-            wl("...")
-            wl("")
-        for i, method in methods:
-            scl = scl_prefix + [d.ServiceDescriptorProto.METHOD_FIELD_NUMBER, i]
-
-            wl("{}: {}", method.name, f"{_build_typevar_name(service.name, method.name)}")
-            self._write_comments(scl)
-            wl("")
-
     def write_grpc_services(
         self,
         services: Iterable[d.ServiceDescriptorProto],
@@ -906,58 +871,60 @@ def write_grpc_services(
         wl = self._write_line
         wl("GRPC_GENERATED_VERSION: str")
         wl("GRPC_VERSION: str")
+        wl("")
         for i, service in enumerate(services):
             if service.name in PYTHON_RESERVED:
                 continue
 
             scl = scl_prefix + [i]
 
-            # Type vars
-            self.write_grpc_type_vars(service)
+            class_name = f"{service.name}Stub"
+            async_class_alias = f"{service.name}AsyncStub"
 
             # The stub client
             if service.options.deprecated:
                 self._write_deprecation_warning(
                     scl + [d.ServiceDescriptorProto.OPTIONS_FIELD_NUMBER] + [d.ServiceOptions.DEPRECATED_FIELD_NUMBER],
                     "This stub has been marked as deprecated using proto service options.",
                 )
-            class_name = f"{service.name}Stub"
             wl(
-                "class {}({}[{}]):",
+                "class {}:",
                 class_name,
-                self._import("typing", "Generic"),
-                ", ".join(f"{_build_typevar_name(service.name, method.name)}" for method in service.method),
             )
             with self._indent():
                 if self._write_comments(scl):
                     wl("")
-
                 # Write sync overload
                 wl("@{}", self._import("typing", "overload"))
-                wl("def __init__(self: {}[", class_name)
-                self.write_self_types(service, False)
                 wl(
-                    "], channel: {}) -> None: ...",
+                    "def __new__(cls, channel: {}) -> {}: ...",
                     self._import("grpc", "Channel"),
+                    class_name,
                 )
-                wl("")
 
                 # Write async overload
                 wl("@{}", self._import("typing", "overload"))
-                wl("def __init__(self: {}[", class_name)
-                self.write_self_types(service, True)
                 wl(
-                    "], channel: {}) -> None: ...",
+                    "def __new__(cls, channel: {}) -> {}: ...",
                     self._import("grpc.aio", "Channel"),
+                    async_class_alias,
                 )
+                self.write_grpc_stub_methods(service, scl, is_async=False)
                 wl("")
 
-                self.write_grpc_stub_methods(service, scl)
-
-            # Write AsyncStub alias
-            wl("{}AsyncStub: {} = {}[", service.name, self._import("typing_extensions", "TypeAlias"), class_name)
-            self.write_self_types(service, True)
-            wl("]")
+            # Write AsyncStub
+            if service.options.deprecated:
+                self._write_deprecation_warning(
+                    scl + [d.ServiceDescriptorProto.OPTIONS_FIELD_NUMBER] + [d.ServiceOptions.DEPRECATED_FIELD_NUMBER],
+                    "This stub has been marked as deprecated using proto service options.",
+                )
+            wl("@{}", self._import("typing", "type_check_only"))
+            wl("class {}({}):", async_class_alias, class_name)
+            with self._indent():
+                if self._write_comments(scl):
+                    wl("")
+                wl("def __init__(self, channel: {}) -> None: ...", self._import("grpc.aio", "Channel"))
+                self.write_grpc_stub_methods(service, scl, is_async=True, ignore_assignment_errors=True)
             wl("")
 
             # The service definition interface