async def _read_messages(self):
"""Read one message. Wire format is 4-byte size, then data."""
while True:
header = await self._conn_read.read(4)
if len(header) == 0:
# Connection closed
return
msg_len = struct.unpack("!I", header)[0]
msg_bytes = await self._conn_read.read(msg_len)
msg = Message()
msg.ParseFromString(msg_bytes)
yield msg
async def reply_to(self, req, data=None):
"""
Send a reply to the request req, with optional data in the reply.
:param Request req: the original request
:param bytes data: optional data to put in the reply
"""
reply = Reply()
reply.id = req.id
if data:
reply.data = data
msg = Message()
msg.type = Message.Reply
msg.content = reply.SerializeToString()
await self.send_message(msg)
def reply_to(self, req, data=None):
"""
Send a reply to the request req, with optional data in the reply.
:param Request req: the original request
:param bytes data: optional data to put in the reply
"""
reply = Reply()
reply.id = req.id
if data:
reply.data = data
msg = Message()
msg.type = Message.Reply
msg.content = reply.SerializeToString()
self.send_message(msg)
async def request(self,
method,
service,
*,
identification=None,
data=None):
"""
Send a ``request`` message.
:param str method: The name of the method.
:param str service: The name of the service.
:return: The id of the message that was sent. Used for tracking the
reply.
:rtype: int
"""
logger.info("[Request] %s/%s.%s(%s)", service, identification, method,
data)
request = Request(service_name=service, method=method)
if identification is not None:
request.service_identification = identification
if data:
request.data = json.dumps(data).encode()
request.id = self._request_seq_id
self._request_seq_id += 1
message = Message(type=Message.Request,
content=request.SerializeToString())
# Create a future that will hold the reply in the result
request_future = asyncio.Future()
self._requests_in_flight[request.id] = request, request_future
await self.send_message(message)
return await request_future
async def reply_error_to(self, req, error_type, what=None):
"""
Send an error reply to the request ``req``.
:param Request req: The original request.
:param Reply.Error error_type: An error code.
:param bytes what: an error message.
"""
error = Reply.Error()
error.type = error_type
if what is not None:
error.what = what
reply = Reply(id=req.id, error=error)
msg = Message()
msg.type = Message.Reply
msg.content = reply.SerializeToString()
await self.send_message(msg)
def reply_error_to(self, req, error_type, what=None):
"""
Send an error reply to the request ``req``.
:param Request req: The original request.
:param Reply.Error error_type: An error code.
:param bytes what: an error message.
"""
error = Reply.Error()
error.type = error_type
if what is not None:
error.what = what
reply = Reply(id=req.id, error=error)
msg = Message()
msg.type = Message.Reply
msg.content = reply.SerializeToString()
self.send_message(msg)
async def subscribe(self, event: str, cb=None) -> None:
"""
Subscribes to an event.
:param str event: The name of the event.
:param cb func or coro: Callback when the event is received.
"""
logger.info("[Subscribe] %s", event)
subscribe = Subscribe(event=event)
message = Message(type=Message.Subscribe,
content=subscribe.SerializeToString())
if cb is not None:
self._subscribes[event].append(cb)
await self.send_message(message)
def publish(self, event, *args, **kwargs):
"""
Send a ``publish`` message.
:param event str: The event name.
:param **kwargs: Optional key=value data sent with the event.
"""
if args and kwargs:
raise TypeError("Cannot send a publish with both args and kwargs")
data = args or kwargs
logger.info("[Publish] %s(%s)", event, data)
publish = Publish(event=event)
try:
publish.data = json.dumps(data).encode()
except TypeError:
logging.error("Could not serialize publish data: %s", data)
raise
message = Message(type=Message.Publish,
content=publish.SerializeToString())
asyncio.create_task(self.send_message(message))
