Skip to content

Commit

Permalink
chore: avoid type checks in error wrapper
Browse files Browse the repository at this point in the history
  • Loading branch information
daniel-sanche authored Sep 7, 2023
1 parent 0cc03bd commit c97a636
Showing 1 changed file with 11 additions and 18 deletions.
29 changes: 11 additions & 18 deletions google/api_core/grpc_helpers_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,8 +149,6 @@ class _WrappedStreamStreamCall(

def _wrap_unary_errors(callable_):
"""Map errors for Unary-Unary async callables."""
grpc_helpers._patch_callable_name(callable_)

@functools.wraps(callable_)
def error_remapped_callable(*args, **kwargs):
call = callable_(*args, **kwargs)
Expand All @@ -159,25 +157,13 @@ def error_remapped_callable(*args, **kwargs):
return error_remapped_callable


def _wrap_stream_errors(callable_):
def _wrap_stream_errors(callable_, wrapper_type):
"""Map errors for streaming RPC async callables."""
grpc_helpers._patch_callable_name(callable_)

@functools.wraps(callable_)
async def error_remapped_callable(*args, **kwargs):
call = callable_(*args, **kwargs)

if isinstance(call, aio.UnaryStreamCall):
call = _WrappedUnaryStreamCall().with_call(call)
elif isinstance(call, aio.StreamUnaryCall):
call = _WrappedStreamUnaryCall().with_call(call)
elif isinstance(call, aio.StreamStreamCall):
call = _WrappedStreamStreamCall().with_call(call)
else:
raise TypeError("Unexpected type of call %s" % type(call))

await call.wait_for_connection()
return call
return wrapper_type().with_call(call)

return error_remapped_callable

Expand All @@ -197,10 +183,17 @@ def wrap_errors(callable_):
Returns: Callable: The wrapped gRPC callable.
"""
grpc_helpers._patch_callable_name(callable_)
if isinstance(callable_, aio.UnaryUnaryMultiCallable):
return _wrap_unary_errors(callable_)
else:
return _wrap_stream_errors(callable_)
if isinstance(callable_, aio.UnaryStreamMultiCallable):
return _wrap_stream_errors(callable_, _WrappedUnaryStreamCall)
elif isinstance(callable_, aio.StreamUnaryMultiCallable):
return _wrap_stream_errors(callable_, _WrappedStreamUnaryCall)
elif isinstance(callable_, aio.StreamStreamMultiCallable):
return _wrap_stream_errors(callable_, _WrappedStreamStreamCall)
else:
raise TypeError("Unexpected type of callable: {}".format(type(callable_)))


def create_channel(
Expand Down

0 comments on commit c97a636

Please sign in to comment.