daniel-sanche · Nov 17, 2023 · Nov 17, 2023 · Nov 17, 2023 · Nov 17, 2023 · Nov 17, 2023
diff --git a/google/api_core/grpc_helpers.py b/google/api_core/grpc_helpers.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Helpers for :mod:`grpc`."""
+from typing import Generic, TypeVar, Iterator
 
 import collections
 import functools
@@ -54,6 +55,9 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+# denotes the type yielded from streaming calls
+S = TypeVar("S")
+
 
 def _patch_callable_name(callable_):
     """Fix-up gRPC callable attributes.
@@ -79,7 +83,7 @@ def error_remapped_callable(*args, **kwargs):
     return error_remapped_callable
 
 
-class _StreamingResponseIterator(grpc.Call):
+class GrpcStream(grpc.Call, Generic[S]):
     def __init__(self, wrapped, prefetch_first_result=True):
         self._wrapped = wrapped
 
@@ -97,11 +101,11 @@ def __init__(self, wrapped, prefetch_first_result=True):
             # ignore stop iteration at this time. This should be handled outside of retry.
             pass
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[S]:
         """This iterator is also an iterable that returns itself."""
         return self
 
-    def __next__(self):
+    def __next__(self) -> S:
         """Get the next response from the stream.
 
         Returns:
@@ -162,9 +166,7 @@ def error_remapped_callable(*args, **kwargs):
             # hidden flag to see if pre-fetching is disabled.
             # https://github.com/googleapis/python-pubsub/issues/93#issuecomment-630762257
             prefetch_first = getattr(callable_, "_prefetch_first_result_", True)
-            return _StreamingResponseIterator(
-                result, prefetch_first_result=prefetch_first
-            )
+            return GrpcStream(result, prefetch_first_result=prefetch_first)
         except grpc.RpcError as exc:
             raise exceptions.from_grpc_error(exc) from exc
 

diff --git a/google/api_core/grpc_helpers_async.py b/google/api_core/grpc_helpers_async.py
@@ -21,6 +21,8 @@
 import asyncio
 import functools
 
+from typing import Generic, Iterator, AsyncGenerator, TypeVar
+
 import grpc
 from grpc import aio
 
@@ -31,6 +33,11 @@
 # automatic patching for us. But that means the overhead of creating an
 # extra Python function spreads to every single send and receive.
 
+# denotes the type returned from unary calls
+U = TypeVar("U")
+# denotes the type yielded from streaming calls
+S = TypeVar("S")
+
 
 class _WrappedCall(aio.Call):
     def __init__(self):
@@ -75,26 +82,26 @@ async def wait_for_connection(self):
             raise exceptions.from_grpc_error(rpc_error) from rpc_error
 
 
-class _WrappedUnaryResponseMixin(_WrappedCall):
-    def __await__(self):
+class _WrappedUnaryResponseMixin(Generic[U], _WrappedCall):
+    def __await__(self) -> Iterator[U]:
         try:
             response = yield from self._call.__await__()
             return response
         except grpc.RpcError as rpc_error:
             raise exceptions.from_grpc_error(rpc_error) from rpc_error
 
 
-class _WrappedStreamResponseMixin(_WrappedCall):
+class _WrappedStreamResponseMixin(Generic[S], _WrappedCall):
     def __init__(self):
         self._wrapped_async_generator = None
 
-    async def read(self):
+    async def read(self) -> S:
         try:
             return await self._call.read()
         except grpc.RpcError as rpc_error:
             raise exceptions.from_grpc_error(rpc_error) from rpc_error
 
-    async def _wrapped_aiter(self):
+    async def _wrapped_aiter(self) -> AsyncGenerator[S, None]:
         try:
             # NOTE(lidiz) coverage doesn't understand the exception raised from
             # __anext__ method. It is covered by test case:
@@ -104,7 +111,7 @@ async def _wrapped_aiter(self):
         except grpc.RpcError as rpc_error:
             raise exceptions.from_grpc_error(rpc_error) from rpc_error
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncGenerator[S, None]:
         if not self._wrapped_async_generator:
             self._wrapped_async_generator = self._wrapped_aiter()
         return self._wrapped_async_generator
@@ -127,26 +134,32 @@ async def done_writing(self):
 # NOTE(lidiz) Implementing each individual class separately, so we don't
 # expose any API that should not be seen. E.g., __aiter__ in unary-unary
 # RPC, or __await__ in stream-stream RPC.
-class _WrappedUnaryUnaryCall(_WrappedUnaryResponseMixin, aio.UnaryUnaryCall):
+class _WrappedUnaryUnaryCall(_WrappedUnaryResponseMixin[U], aio.UnaryUnaryCall):
     """Wrapped UnaryUnaryCall to map exceptions."""
 
 
-class _WrappedUnaryStreamCall(_WrappedStreamResponseMixin, aio.UnaryStreamCall):
+class _WrappedUnaryStreamCall(_WrappedStreamResponseMixin[S], aio.UnaryStreamCall):
     """Wrapped UnaryStreamCall to map exceptions."""
 
 
 class _WrappedStreamUnaryCall(
-    _WrappedUnaryResponseMixin, _WrappedStreamRequestMixin, aio.StreamUnaryCall
+    _WrappedUnaryResponseMixin[S], _WrappedStreamRequestMixin, aio.StreamUnaryCall
 ):
     """Wrapped StreamUnaryCall to map exceptions."""
 
 
 class _WrappedStreamStreamCall(
-    _WrappedStreamRequestMixin, _WrappedStreamResponseMixin, aio.StreamStreamCall
+    _WrappedStreamRequestMixin, _WrappedStreamResponseMixin[S], aio.StreamStreamCall
 ):
     """Wrapped StreamStreamCall to map exceptions."""
 
 
+# public type alias denoting the return type of streaming gapic calls
+GrpcAsyncStream = _WrappedStreamResponseMixin[S]
+# public type alias denoting the return type of unary gapic calls
+AwaitableGrpcCall = _WrappedUnaryResponseMixin[U]
+
+
 def _wrap_unary_errors(callable_):
     """Map errors for Unary-Unary async callables."""
     grpc_helpers._patch_callable_name(callable_)

diff --git a/tests/asyncio/test_grpc_helpers_async.py b/tests/asyncio/test_grpc_helpers_async.py
@@ -266,6 +266,28 @@ def test_wrap_errors_non_streaming(wrap_unary_errors):
     wrap_unary_errors.assert_called_once_with(callable_)
 
 
+def test_grpc_async_stream():
+    """
+    GrpcAsyncStream type be both an AsyncIterator and a grpc.aio.Call.
+    """
+    instance = grpc_helpers_async.GrpcAsyncStream[int]()
+    assert isinstance(instance, grpc.aio.Call)
+    # should implement __aiter__ and __anext__
+    assert hasattr(instance, "__aiter__")
+    it = instance.__aiter__()
+    assert hasattr(it, "__anext__")
+
+
+def test_awaitable_grpc_call():
+    """
+    AwaitableGrpcCall type should be an Awaitable and a grpc.aio.Call.
+    """
+    instance = grpc_helpers_async.AwaitableGrpcCall[int]()
+    assert isinstance(instance, grpc.aio.Call)
+    # should implement __await__
+    assert hasattr(instance, "__await__")
+
+
 @mock.patch("google.api_core.grpc_helpers_async._wrap_stream_errors")
 def test_wrap_errors_streaming(wrap_stream_errors):
     callable_ = mock.create_autospec(aio.UnaryStreamMultiCallable)

diff --git a/tests/unit/test_grpc_helpers.py b/tests/unit/test_grpc_helpers.py
@@ -73,14 +73,14 @@ def test_wrap_unary_errors():
     assert exc_info.value.response == grpc_error
 
 
-class Test_StreamingResponseIterator:
+class TestGrpcStream:
     @staticmethod
     def _make_wrapped(*items):
         return iter(items)
 
     @staticmethod
     def _make_one(wrapped, **kw):
-        return grpc_helpers._StreamingResponseIterator(wrapped, **kw)
+        return grpc_helpers.GrpcStream(wrapped, **kw)
 
     def test_ctor_defaults(self):
         wrapped = self._make_wrapped("a", "b", "c")