pw_rpc/py/pw_rpc/callback_client.py - pigweed/pigweed - Git at Google

 # Copyright 2020 The Pigweed Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
 # use this file except in compliance with the License. You may obtain a copy of
 # the License at
 #
 #     https://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations under
 # the License.
 """Defines a callback-based RPC ClientImpl to use with pw_rpc.client.Client.

 callback_client.Impl supports invoking RPCs synchronously or asynchronously.
 Asynchronous invocations use a callback.

 Synchronous invocations look like a function call:

   status, response = client.channel(1).call.MyServer.MyUnary(some_field=123)

   # Streaming calls return an iterable of responses
   for reply in client.channel(1).call.MyService.MyServerStreaming(request):
       pass

 Asynchronous invocations pass a callback in addition to the request. The
 callback must be a callable that accepts a status and a payload, either of
 which may be None. The Status is only set when the RPC is completed.

   callback = lambda status, payload: print('Response:', status, payload)

   call = client.channel(1).call.MyServer.MyUnary.invoke(
       callback, some_field=123)

   call = client.channel(1).call.MyService.MyServerStreaming.invoke(
       callback, request):

 When invoking a method, requests may be provided as a message object or as
 kwargs for the message fields (but not both).
 """

 import inspect
 import logging
 import queue
 from typing import Any, Callable, Optional, Tuple

 from pw_status import Status

 from pw_rpc import client
 from pw_rpc.descriptors import Channel, Method, Service

 _LOG = logging.getLogger(__name__)

 Callback = Callable[[client.PendingRpc, Optional[Status], Any], Any]


 class _MethodClient:
     """A method that can be invoked for a particular channel."""
     def __init__(self, client_impl: 'Impl', rpcs: client.PendingRpcs,
                  channel: Channel, method: Method):
         self._impl = client_impl
         self._rpcs = rpcs
         self._rpc = client.PendingRpc(channel, method.service, method)

     @property
     def channel(self) -> Channel:
         return self._rpc.channel

     @property
     def method(self) -> Method:
         return self._rpc.method

     @property
     def service(self) -> Service:
         return self._rpc.service

     def invoke(self, callback: Callback, _request=None, **request_fields):
         """Invokes an RPC with a callback."""
         self._rpcs.send_request(self._rpc,
                                 self.method.get_request(
                                     _request, request_fields),
                                 callback,
                                 override_pending=False)
         return _AsyncCall(self._rpcs, self._rpc)

     def reinvoke(self, callback: Callback, _request=None, **request_fields):
         """Invokes an RPC with a callback, overriding any pending requests."""
         self._rpcs.send_request(self._rpc,
                                 self.method.get_request(
                                     _request, request_fields),
                                 callback,
                                 override_pending=True)
         return _AsyncCall(self._rpcs, self._rpc)

     def __repr__(self) -> str:
         return repr(self.method)


 class _AsyncCall:
     """Represents an ongoing callback-based call."""

     # TODO(hepler): Consider alternatives (futures) and/or expand functionality.

     def __init__(self, rpcs: client.PendingRpcs, rpc: client.PendingRpc):
         self.rpc = rpc
         self._rpcs = rpcs

     def cancel(self) -> bool:
         return self._rpcs.send_cancel(self.rpc)

     def __enter__(self) -> '_AsyncCall':
         return self

     def __exit__(self, exc_type, exc_value, traceback) -> None:
         self.cancel()


 class _StreamingResponses:
     """Used to iterate over a queue.SimpleQueue."""
     def __init__(self, responses: queue.SimpleQueue):
         self._queue = responses
         self.status: Optional[Status] = None

     def get(self, *, block: bool = True, timeout_s: float = None):
         while True:
             self.status, response = self._queue.get(block, timeout_s)
             if self.status is not None:
                 return

             yield response

     def __iter__(self):
         return self.get()


 def _method_client_docstring(method: Method) -> str:
     return f'''\
 Class that invokes the {method.full_name} {method.type.sentence_name()} RPC.

 Calling this directly invokes the RPC synchronously. The RPC can be invoked
 asynchronously using the invoke or reinvoke methods.
 '''


 def _function_docstring(method: Method) -> str:
     return f'''\
 Invokes the {method.full_name} {method.type.sentence_name()} RPC.

 This function accepts either a request protobuf as a single positional argument
 or the request fields as keyword arguments.
 '''


 def _update_function_signature(method: Method, function: Callable) -> None:
     """Updates the name, docstring, and parameters to match a method."""
     function.__name__ = method.full_name
     function.__doc__ = _function_docstring(method)

     # In order to have good tab completion and help messages, update the
     # function signature to accept only keyword arguments for the proto message
     # fields. This doesn't actually change the function signature -- it just
     # updates how it appears when inspected.
     sig = inspect.signature(function)

     params = [next(iter(sig.parameters.values()))]  # Get the "self" parameter
     params += method.request_parameters()
     function.__signature__ = sig.replace(  # type: ignore[attr-defined]
         parameters=params)


 def unary_method_client(client_impl: 'Impl', rpcs: client.PendingRpcs,
                         channel: Channel, method: Method) -> _MethodClient:
     """Creates an object used to call a unary method."""
     def call(self,
              _rpc_request_proto=None,
              **request_fields) -> Tuple[Status, Any]:
         responses: queue.SimpleQueue = queue.SimpleQueue()
         self.reinvoke(
             lambda _, status, payload: responses.put((status, payload)),
             _rpc_request_proto, **request_fields)
         return responses.get()

     _update_function_signature(method, call)

     # The MethodClient class is created dynamically so that the __call__ method
     # can be configured differently for each method.
     method_client_type = type(
         f'{method.name}_UnaryMethodClient', (_MethodClient, ),
         dict(__call__=call, __doc__=_method_client_docstring(method)))
     return method_client_type(client_impl, rpcs, channel, method)


 def server_streaming_method_client(client_impl: 'Impl',
                                    rpcs: client.PendingRpcs, channel: Channel,
                                    method: Method):
     """Creates an object used to call a server streaming method."""
     def call(self,
              _rpc_request_proto=None,
              **request_fields) -> _StreamingResponses:
         responses: queue.SimpleQueue = queue.SimpleQueue()
         self.reinvoke(
             lambda _, status, payload: responses.put((status, payload)),
             _rpc_request_proto, **request_fields)
         return _StreamingResponses(responses)

     _update_function_signature(method, call)

     # The MethodClient class is created dynamically so that the __call__ method
     # can be configured differently for each method type.
     method_client_type = type(
         f'{method.name}_ServerStreamingMethodClient', (_MethodClient, ),
         dict(__call__=call, __doc__=_method_client_docstring(method)))
     return method_client_type(client_impl, rpcs, channel, method)


 class ClientStreamingMethodClient(_MethodClient):
     def __call__(self):
         raise NotImplementedError

     def invoke(self, callback: Callback, _request=None, **request_fields):
         raise NotImplementedError


 class BidirectionalStreamingMethodClient(_MethodClient):
     def __call__(self):
         raise NotImplementedError

     def invoke(self, callback: Callback, _request=None, **request_fields):
         raise NotImplementedError


 class Impl(client.ClientImpl):
     """Callback-based client.ClientImpl."""
     def method_client(self, rpcs: client.PendingRpcs, channel: Channel,
                       method: Method) -> _MethodClient:
         """Returns an object that invokes a method using the given chanel."""

         if method.type is Method.Type.UNARY:
             return unary_method_client(self, rpcs, channel, method)

         if method.type is Method.Type.SERVER_STREAMING:
             return server_streaming_method_client(self, rpcs, channel, method)

         if method.type is Method.Type.CLIENT_STREAMING:
             return ClientStreamingMethodClient(self, rpcs, channel, method)

         if method.type is Method.Type.BIDIRECTIONAL_STREAMING:
             return BidirectionalStreamingMethodClient(self, rpcs, channel,
                                                       method)

         raise AssertionError(f'Unknown method type {method.type}')

     def process_response(self,
                          rpcs: client.PendingRpcs,
                          rpc: client.PendingRpc,
                          context,
                          status: Optional[Status],
                          payload,
                          *,
                          args: tuple = (),
                          kwargs: dict = None) -> None:
         """Invokes the callback associated with this RPC.

         Any additional positional and keyword args passed through
         Client.process_packet are forwarded to the callback.
         """
         if kwargs is None:
             kwargs = {}

         try:
             context(rpc, status, payload, *args, **kwargs)
         except:  # pylint: disable=bare-except
             rpcs.send_cancel(rpc)
             _LOG.exception('Callback %s for %s raised exception', context, rpc)
	# Copyright 2020 The Pigweed Authors
	#
	# Licensed under the Apache License, Version 2.0 (the "License"); you may not
	# use this file except in compliance with the License. You may obtain a copy of
	# the License at
	#
	# https://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
	# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
	# License for the specific language governing permissions and limitations under
	# the License.
	"""Defines a callback-based RPC ClientImpl to use with pw_rpc.client.Client.

	callback_client.Impl supports invoking RPCs synchronously or asynchronously.
	Asynchronous invocations use a callback.

	Synchronous invocations look like a function call:

	status, response = client.channel(1).call.MyServer.MyUnary(some_field=123)

	# Streaming calls return an iterable of responses
	for reply in client.channel(1).call.MyService.MyServerStreaming(request):
	pass

	Asynchronous invocations pass a callback in addition to the request. The
	callback must be a callable that accepts a status and a payload, either of
	which may be None. The Status is only set when the RPC is completed.

	callback = lambda status, payload: print('Response:', status, payload)

	call = client.channel(1).call.MyServer.MyUnary.invoke(
	callback, some_field=123)

	call = client.channel(1).call.MyService.MyServerStreaming.invoke(
	callback, request):

	When invoking a method, requests may be provided as a message object or as
	kwargs for the message fields (but not both).
	"""

	import inspect
	import logging
	import queue
	from typing import Any, Callable, Optional, Tuple

	from pw_status import Status

	from pw_rpc import client
	from pw_rpc.descriptors import Channel, Method, Service

	_LOG = logging.getLogger(__name__)

	Callback = Callable[[client.PendingRpc, Optional[Status], Any], Any]


	class _MethodClient:
	"""A method that can be invoked for a particular channel."""
	def __init__(self, client_impl: 'Impl', rpcs: client.PendingRpcs,
	channel: Channel, method: Method):
	self._impl = client_impl
	self._rpcs = rpcs
	self._rpc = client.PendingRpc(channel, method.service, method)

	@property
	def channel(self) -> Channel:
	return self._rpc.channel

	@property
	def method(self) -> Method:
	return self._rpc.method

	@property
	def service(self) -> Service:
	return self._rpc.service

	def invoke(self, callback: Callback, _request=None, **request_fields):
	"""Invokes an RPC with a callback."""
	self._rpcs.send_request(self._rpc,
	self.method.get_request(
	_request, request_fields),
	callback,
	override_pending=False)
	return _AsyncCall(self._rpcs, self._rpc)

	def reinvoke(self, callback: Callback, _request=None, **request_fields):
	"""Invokes an RPC with a callback, overriding any pending requests."""
	self._rpcs.send_request(self._rpc,
	self.method.get_request(
	_request, request_fields),
	callback,
	override_pending=True)
	return _AsyncCall(self._rpcs, self._rpc)

	def __repr__(self) -> str:
	return repr(self.method)


	class _AsyncCall:
	"""Represents an ongoing callback-based call."""

	# TODO(hepler): Consider alternatives (futures) and/or expand functionality.

	def __init__(self, rpcs: client.PendingRpcs, rpc: client.PendingRpc):
	self.rpc = rpc
	self._rpcs = rpcs

	def cancel(self) -> bool:
	return self._rpcs.send_cancel(self.rpc)

	def __enter__(self) -> '_AsyncCall':
	return self

	def __exit__(self, exc_type, exc_value, traceback) -> None:
	self.cancel()


	class _StreamingResponses:
	"""Used to iterate over a queue.SimpleQueue."""
	def __init__(self, responses: queue.SimpleQueue):
	self._queue = responses
	self.status: Optional[Status] = None

	def get(self, *, block: bool = True, timeout_s: float = None):
	while True:
	self.status, response = self._queue.get(block, timeout_s)
	if self.status is not None:
	return

	yield response

	def __iter__(self):
	return self.get()


	def _method_client_docstring(method: Method) -> str:
	return f'''\
	Class that invokes the {method.full_name} {method.type.sentence_name()} RPC.

	Calling this directly invokes the RPC synchronously. The RPC can be invoked
	asynchronously using the invoke or reinvoke methods.
	'''


	def _function_docstring(method: Method) -> str:
	return f'''\
	Invokes the {method.full_name} {method.type.sentence_name()} RPC.

	This function accepts either a request protobuf as a single positional argument
	or the request fields as keyword arguments.
	'''


	def _update_function_signature(method: Method, function: Callable) -> None:
	"""Updates the name, docstring, and parameters to match a method."""
	function.__name__ = method.full_name
	function.__doc__ = _function_docstring(method)

	# In order to have good tab completion and help messages, update the
	# function signature to accept only keyword arguments for the proto message
	# fields. This doesn't actually change the function signature -- it just
	# updates how it appears when inspected.
	sig = inspect.signature(function)

	params = [next(iter(sig.parameters.values()))] # Get the "self" parameter
	params += method.request_parameters()
	function.__signature__ = sig.replace( # type: ignore[attr-defined]
	parameters=params)


	def unary_method_client(client_impl: 'Impl', rpcs: client.PendingRpcs,
	channel: Channel, method: Method) -> _MethodClient:
	"""Creates an object used to call a unary method."""
	def call(self,
	_rpc_request_proto=None,
	**request_fields) -> Tuple[Status, Any]:
	responses: queue.SimpleQueue = queue.SimpleQueue()
	self.reinvoke(
	lambda _, status, payload: responses.put((status, payload)),
	_rpc_request_proto, **request_fields)
	return responses.get()

	_update_function_signature(method, call)

	# The MethodClient class is created dynamically so that the __call__ method
	# can be configured differently for each method.
	method_client_type = type(
	f'{method.name}_UnaryMethodClient', (_MethodClient, ),
	dict(__call__=call, __doc__=_method_client_docstring(method)))
	return method_client_type(client_impl, rpcs, channel, method)


	def server_streaming_method_client(client_impl: 'Impl',
	rpcs: client.PendingRpcs, channel: Channel,
	method: Method):
	"""Creates an object used to call a server streaming method."""
	def call(self,
	_rpc_request_proto=None,
	**request_fields) -> _StreamingResponses:
	responses: queue.SimpleQueue = queue.SimpleQueue()
	self.reinvoke(
	lambda _, status, payload: responses.put((status, payload)),
	_rpc_request_proto, **request_fields)
	return _StreamingResponses(responses)

	_update_function_signature(method, call)

	# The MethodClient class is created dynamically so that the __call__ method
	# can be configured differently for each method type.
	method_client_type = type(
	f'{method.name}_ServerStreamingMethodClient', (_MethodClient, ),
	dict(__call__=call, __doc__=_method_client_docstring(method)))
	return method_client_type(client_impl, rpcs, channel, method)


	class ClientStreamingMethodClient(_MethodClient):
	def __call__(self):
	raise NotImplementedError

	def invoke(self, callback: Callback, _request=None, **request_fields):
	raise NotImplementedError


	class BidirectionalStreamingMethodClient(_MethodClient):
	def __call__(self):
	raise NotImplementedError

	def invoke(self, callback: Callback, _request=None, **request_fields):
	raise NotImplementedError


	class Impl(client.ClientImpl):
	"""Callback-based client.ClientImpl."""
	def method_client(self, rpcs: client.PendingRpcs, channel: Channel,
	method: Method) -> _MethodClient:
	"""Returns an object that invokes a method using the given chanel."""

	if method.type is Method.Type.UNARY:
	return unary_method_client(self, rpcs, channel, method)

	if method.type is Method.Type.SERVER_STREAMING:
	return server_streaming_method_client(self, rpcs, channel, method)

	if method.type is Method.Type.CLIENT_STREAMING:
	return ClientStreamingMethodClient(self, rpcs, channel, method)

	if method.type is Method.Type.BIDIRECTIONAL_STREAMING:
	return BidirectionalStreamingMethodClient(self, rpcs, channel,
	method)

	raise AssertionError(f'Unknown method type {method.type}')

	def process_response(self,
	rpcs: client.PendingRpcs,
	rpc: client.PendingRpc,
	context,
	status: Optional[Status],
	payload,
	*,
	args: tuple = (),
	kwargs: dict = None) -> None:
	"""Invokes the callback associated with this RPC.

	Any additional positional and keyword args passed through
	Client.process_packet are forwarded to the callback.
	"""
	if kwargs is None:
	kwargs = {}

	try:
	context(rpc, status, payload, args, *kwargs)
	except: # pylint: disable=bare-except
	rpcs.send_cancel(rpc)
	_LOG.exception('Callback %s for %s raised exception', context, rpc)