text: import asyncio
  import functools
  from typing import Callable
  from typing import Tuple
  from typing import Union
  
  import grpc
  from grpc import aio
  from grpc.aio._typing import RequestIterableType
  from grpc.aio._typing import RequestType
  from grpc.aio._typing import ResponseIterableType
  from grpc.aio._typing import ResponseType
  
  from .. import trace_utils
  from ... import Pin
  from ... import Span
  from ... import config
  from ...constants import ANALYTICS_SAMPLE_RATE_KEY
  from ...constants import ERROR_MSG
  from ...constants import ERROR_TYPE
  from ...constants import SPAN_MEASURED_KEY
  from ...ext import SpanTypes
  from ...internal.compat import to_unicode
  from ...propagation.http import HTTPPropagator
  from ..grpc import constants
  from ..grpc import utils
  
  
  def create_aio_client_interceptors(pin, host, port):
      # type: (Pin, str, int) -> Tuple[aio.ClientInterceptor, ...]
      return (
          _UnaryUnaryClientInterceptor(pin, host, port),
          _UnaryStreamClientInterceptor(pin, host, port),
          _StreamUnaryClientInterceptor(pin, host, port),
          _StreamStreamClientInterceptor(pin, host, port),
      )
  
  
  def _done_callback(span, code, details):
      # type: (Span, grpc.StatusCode, str) -> Callable[[aio.Call], None]
      def func(call):
          # type: (aio.Call) -> None
          try:
              span._set_str_tag(constants.GRPC_STATUS_CODE_KEY, to_unicode(code))
  
              # Handle server-side error in unary response RPCs
              if code != grpc.StatusCode.OK:
                  _handle_error(span, call, code, details)
          finally:
              span.finish()
  
      return func
  
  
  def _handle_error(span, call, code, details):
      # type: (Span, aio.Call, grpc.StatusCode, str) -> None
      span.error = 1
      span._set_str_tag(ERROR_MSG, details)
      span._set_str_tag(ERROR_TYPE, to_unicode(code))
  
  
  def _handle_rpc_error(span, rpc_error):
      # type: (Span, aio.AioRpcError) -> None
      code = to_unicode(rpc_error.code())
      span.error = 1
      span._set_str_tag(constants.GRPC_STATUS_CODE_KEY, code)
      span._set_str_tag(ERROR_MSG, rpc_error.details())
      span._set_str_tag(ERROR_TYPE, code)
      span.finish()
  
  
  async def _handle_cancelled_error(call, span):
      # type: (aio.Call, Span) -> None
      code = to_unicode(await call.code())
      span.error = 1
      span._set_str_tag(constants.GRPC_STATUS_CODE_KEY, code)
      span._set_str_tag(ERROR_MSG, await call.details())
      span._set_str_tag(ERROR_TYPE, code)
      span.finish()
  
  
  class _ClientInterceptor:
      def __init__(self, pin, host, port):
          # type: (Pin, str, int) -> None
          self._pin = pin
          self._host = host
          self._port = port
  
      def _intercept_client_call(self, method_kind, client_call_details):
          # type: (str, aio.ClientCallDetails) -> Tuple[Span, aio.ClientCallDetails]
          tracer = self._pin.tracer
  
          method_as_str = client_call_details.method.decode()
          span = tracer.trace(
              "grpc",
              span_type=SpanTypes.GRPC,
              service=trace_utils.ext_service(self._pin, config.grpc_aio_client),
              resource=method_as_str,
          )
          span.set_tag(SPAN_MEASURED_KEY)
  
          utils.set_grpc_method_meta(span, method_as_str, method_kind)
          utils.set_grpc_client_meta(span, self._host, self._port)
          span._set_str_tag(constants.GRPC_SPAN_KIND_KEY, constants.GRPC_SPAN_KIND_VALUE_CLIENT)
  
          sample_rate = config.grpc_aio_client.get_analytics_sample_rate()
          if sample_rate is not None:
              span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, sample_rate)
  
          # inject tags from pin
          if self._pin.tags:
              span.set_tags(self._pin.tags)
  
          # propagate distributed tracing headers if available
          headers = {}
          if config.grpc_aio_client.distributed_tracing_enabled:
              HTTPPropagator.inject(span.context, headers)
  
          metadata = []
          if client_call_details.metadata is not None:
              metadata = list(client_call_details.metadata)
          metadata.extend(headers.items())
  
          client_call_details = aio.ClientCallDetails(
              client_call_details.method,
              client_call_details.timeout,
              metadata,
              client_call_details.credentials,
              client_call_details.wait_for_ready,
          )
  
          return span, client_call_details
  
      # NOTE: Since this function is executed as an async generator when the RPC is called,
      # `continuation` must be called before the RPC.
      async def _wrap_stream_response(
          self,
          call,  # type: Union[aio.StreamStreamCall, aio.UnaryStreamCall]
          span,  # type: Span
      ):
          # type: (...) -> ResponseIterableType
          try:
              async for response in call:
                  yield response
              code = await call.code()
              details = await call.details()
              # NOTE: The callback is registered after the iteration is done,
              # otherwise `call.code()` and `call.details()` block indefinitely.
              call.add_done_callback(_done_callback(span, code, details))
          except aio.AioRpcError as rpc_error:
              # NOTE: We can also handle the error in done callbacks,
              # but reuse this error handling function used in unary response RPCs.
              _handle_rpc_error(span, rpc_error)
              raise
          except asyncio.CancelledError:
              # NOTE: We can't handle the cancelled error in done callbacks
              # because they cannot handle awaitable functions.
              await _handle_cancelled_error(call, span)
              raise
  
      # NOTE: `continuation` must be called inside of this function to catch exceptions.
      async def _wrap_unary_response(
          self,
          continuation,  # type: Callable[[], Union[aio.StreamUnaryCall, aio.UnaryUnaryCall]]
          span,  # type: Span
      ):
          # type: (...) -> Union[aio.StreamUnaryCall, aio.UnaryUnaryCall]
          try:
              call = await continuation()
              code = await call.code()
              details = await call.details()
              # NOTE: As both `code` and `details` are available after the RPC is done (= we get `call` object),
              # and we can't call awaitable functions inside the non-async callback,
              # there is no other way but to register the callback here.
              call.add_done_callback(_done_callback(span, code, details))
              return call
          except aio.AioRpcError as rpc_error:
              # NOTE: `AioRpcError` is raised in `await continuation(...)`
              # and `call` object is not assigned yet in that case.
              # So we can't handle the error in done callbacks.
              _handle_rpc_error(span, rpc_error)
              raise
  
  
  class _UnaryUnaryClientInterceptor(aio.UnaryUnaryClientInterceptor, _ClientInterceptor):
      async def intercept_unary_unary(
          self,
          continuation,  # type: Callable[[aio.ClientCallDetails, RequestType], aio.UnaryUnaryCall]
          client_call_details,  # type: aio.ClientCallDetails
          request,  # type: RequestType
      ):
          # type: (...) -> Union[aio.UnaryUnaryCall, ResponseType]
          span, client_call_details = self._intercept_client_call(
              constants.GRPC_METHOD_KIND_UNARY,
              client_call_details,
          )
          continuation_with_args = functools.partial(continuation, client_call_details, request)
          return await self._wrap_unary_response(continuation_with_args, span)
  
  
  class _UnaryStreamClientInterceptor(aio.UnaryStreamClientInterceptor, _ClientInterceptor):
      async def intercept_unary_stream(
          self,
          continuation,  # type: Callable[[aio.ClientCallDetails, RequestType], aio.UnaryStreamCall]
          client_call_details,  # type: aio.ClientCallDetails
          request,  # type: RequestType
      ):
          # type: (...) -> Union[aio.UnaryStreamCall, ResponseIterableType]
          span, client_call_details = self._intercept_client_call(
              constants.GRPC_METHOD_KIND_SERVER_STREAMING,
              client_call_details,
          )
          call = await continuation(client_call_details, request)
          return self._wrap_stream_response(call, span)
  
  
  class _StreamUnaryClientInterceptor(aio.StreamUnaryClientInterceptor, _ClientInterceptor):
      async def intercept_stream_unary(
          self,
          continuation,  # type: Callable[[aio.ClientCallDetails, RequestType], aio.StreamUnaryCall]
          client_call_details,  # type: aio.ClientCallDetails
          request_iterator,  # type: RequestIterableType
      ):
          # type: (...) -> aio.StreamUnaryCall
          span, client_call_details = self._intercept_client_call(
              constants.GRPC_METHOD_KIND_CLIENT_STREAMING,
              client_call_details,
          )
          continuation_with_args = functools.partial(continuation, client_call_details, request_iterator)
          return await self._wrap_unary_response(continuation_with_args, span)
  
  
  class _StreamStreamClientInterceptor(aio.StreamStreamClientInterceptor, _ClientInterceptor):
      async def intercept_stream_stream(
          self,
          continuation,  # type: Callable[[aio.ClientCallDetails, RequestType], aio.StreamStreamCall]
          client_call_details,  # type: aio.ClientCallDetails
          request_iterator,  # type: RequestIterableType
      ):
          # type: (...) -> Union[aio.StreamStreamCall, ResponseIterableType]
          span, client_call_details = self._intercept_client_call(
              constants.GRPC_METHOD_KIND_BIDI_STREAMING,
              client_call_details,
          )
          call = await continuation(client_call_details, request_iterator)
          return self._wrap_stream_response(call, span)
  