proto.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. # Protocol Buffers - Google's data interchange format
  2. # Copyright 2008 Google Inc. All rights reserved.
  3. #
  4. # Use of this source code is governed by a BSD-style
  5. # license that can be found in the LICENSE file or at
  6. # https://developers.google.com/open-source/licenses/bsd
  7. """Contains the Nextgen Pythonic protobuf APIs."""
  8. import io
  9. from typing import Type, TypeVar
  10. from google.protobuf.internal import decoder
  11. from google.protobuf.internal import encoder
  12. from google.protobuf.message import Message
  13. _MESSAGE = TypeVar('_MESSAGE', bound='Message')
  14. def serialize(message: _MESSAGE, deterministic: bool = None) -> bytes:
  15. """Return the serialized proto.
  16. Args:
  17. message: The proto message to be serialized.
  18. deterministic: If true, requests deterministic serialization
  19. of the protobuf, with predictable ordering of map keys.
  20. Returns:
  21. A binary bytes representation of the message.
  22. """
  23. return message.SerializeToString(deterministic=deterministic)
  24. def parse(message_class: Type[_MESSAGE], payload: bytes) -> _MESSAGE:
  25. """Given a serialized data in binary form, deserialize it into a Message.
  26. Args:
  27. message_class: The message meta class.
  28. payload: A serialized bytes in binary form.
  29. Returns:
  30. A new message deserialized from payload.
  31. """
  32. new_message = message_class()
  33. new_message.ParseFromString(payload)
  34. return new_message
  35. def serialize_length_prefixed(message: _MESSAGE, output: io.BytesIO) -> None:
  36. """Writes the size of the message as a varint and the serialized message.
  37. Writes the size of the message as a varint and then the serialized message.
  38. This allows more data to be written to the output after the message. Use
  39. parse_length_prefixed to parse messages written by this method.
  40. The output stream must be buffered, e.g. using
  41. https://docs.python.org/3/library/io.html#buffered-streams.
  42. Example usage:
  43. out = io.BytesIO()
  44. for msg in message_list:
  45. proto.serialize_length_prefixed(msg, out)
  46. Args:
  47. message: The protocol buffer message that should be serialized.
  48. output: BytesIO or custom buffered IO that data should be written to.
  49. """
  50. size = message.ByteSize()
  51. encoder._VarintEncoder()(output.write, size)
  52. out_size = output.write(serialize(message))
  53. if out_size != size:
  54. raise TypeError(
  55. 'Failed to write complete message (wrote: %d, expected: %d)'
  56. '. Ensure output is using buffered IO.' % (out_size, size)
  57. )
  58. def parse_length_prefixed(
  59. message_class: Type[_MESSAGE], input_bytes: io.BytesIO
  60. ) -> _MESSAGE:
  61. """Parse a message from input_bytes.
  62. Args:
  63. message_class: The protocol buffer message class that parser should parse.
  64. input_bytes: A buffered input.
  65. Example usage:
  66. while True:
  67. msg = proto.parse_length_prefixed(message_class, input_bytes)
  68. if msg is None:
  69. break
  70. ...
  71. Returns:
  72. A parsed message if successful. None if input_bytes is at EOF.
  73. """
  74. size = decoder._DecodeVarint(input_bytes)
  75. if size is None:
  76. # It is the end of buffered input. See example usage in the
  77. # API description.
  78. return None
  79. message = message_class()
  80. if size == 0:
  81. return message
  82. parsed_size = message.ParseFromString(input_bytes.read(size))
  83. if parsed_size != size:
  84. raise ValueError(
  85. 'Truncated message or non-buffered input_bytes: '
  86. 'Expected {0} bytes but only {1} bytes parsed for '
  87. '{2}.'.format(size, parsed_size, message.DESCRIPTOR.name)
  88. )
  89. return message