proto_builder.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. # Protocol Buffers - Google's data interchange format
  2. # Copyright 2008 Google Inc. All rights reserved.
  3. # https://developers.google.com/protocol-buffers/
  4. #
  5. # Redistribution and use in source and binary forms, with or without
  6. # modification, are permitted provided that the following conditions are
  7. # met:
  8. #
  9. # * Redistributions of source code must retain the above copyright
  10. # notice, this list of conditions and the following disclaimer.
  11. # * Redistributions in binary form must reproduce the above
  12. # copyright notice, this list of conditions and the following disclaimer
  13. # in the documentation and/or other materials provided with the
  14. # distribution.
  15. # * Neither the name of Google Inc. nor the names of its
  16. # contributors may be used to endorse or promote products derived from
  17. # this software without specific prior written permission.
  18. #
  19. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  20. # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  21. # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  22. # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  23. # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  24. # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  25. # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  26. # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  27. # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  28. # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  29. # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  30. """Dynamic Protobuf class creator."""
  31. try:
  32. from collections import OrderedDict
  33. except ImportError:
  34. from ordereddict import OrderedDict #PY26
  35. import hashlib
  36. import os
  37. from google.protobuf import descriptor_pb2
  38. from google.protobuf import message_factory
  39. def _GetMessageFromFactory(factory, full_name):
  40. """Get a proto class from the MessageFactory by name.
  41. Args:
  42. factory: a MessageFactory instance.
  43. full_name: str, the fully qualified name of the proto type.
  44. Returns:
  45. A class, for the type identified by full_name.
  46. Raises:
  47. KeyError, if the proto is not found in the factory's descriptor pool.
  48. """
  49. proto_descriptor = factory.pool.FindMessageTypeByName(full_name)
  50. proto_cls = factory.GetPrototype(proto_descriptor)
  51. return proto_cls
  52. def MakeSimpleProtoClass(fields, full_name=None, pool=None):
  53. """Create a Protobuf class whose fields are basic types.
  54. Note: this doesn't validate field names!
  55. Args:
  56. fields: dict of {name: field_type} mappings for each field in the proto. If
  57. this is an OrderedDict the order will be maintained, otherwise the
  58. fields will be sorted by name.
  59. full_name: optional str, the fully-qualified name of the proto type.
  60. pool: optional DescriptorPool instance.
  61. Returns:
  62. a class, the new protobuf class with a FileDescriptor.
  63. """
  64. factory = message_factory.MessageFactory(pool=pool)
  65. if full_name is not None:
  66. try:
  67. proto_cls = _GetMessageFromFactory(factory, full_name)
  68. return proto_cls
  69. except KeyError:
  70. # The factory's DescriptorPool doesn't know about this class yet.
  71. pass
  72. # Get a list of (name, field_type) tuples from the fields dict. If fields was
  73. # an OrderedDict we keep the order, but otherwise we sort the field to ensure
  74. # consistent ordering.
  75. field_items = fields.items()
  76. if not isinstance(fields, OrderedDict):
  77. field_items = sorted(field_items)
  78. # Use a consistent file name that is unlikely to conflict with any imported
  79. # proto files.
  80. fields_hash = hashlib.sha1()
  81. for f_name, f_type in field_items:
  82. fields_hash.update(f_name.encode('utf-8'))
  83. fields_hash.update(str(f_type).encode('utf-8'))
  84. proto_file_name = fields_hash.hexdigest() + '.proto'
  85. # If the proto is anonymous, use the same hash to name it.
  86. if full_name is None:
  87. full_name = ('net.proto2.python.public.proto_builder.AnonymousProto_' +
  88. fields_hash.hexdigest())
  89. try:
  90. proto_cls = _GetMessageFromFactory(factory, full_name)
  91. return proto_cls
  92. except KeyError:
  93. # The factory's DescriptorPool doesn't know about this class yet.
  94. pass
  95. # This is the first time we see this proto: add a new descriptor to the pool.
  96. factory.pool.Add(
  97. _MakeFileDescriptorProto(proto_file_name, full_name, field_items))
  98. return _GetMessageFromFactory(factory, full_name)
  99. def _MakeFileDescriptorProto(proto_file_name, full_name, field_items):
  100. """Populate FileDescriptorProto for MessageFactory's DescriptorPool."""
  101. package, name = full_name.rsplit('.', 1)
  102. file_proto = descriptor_pb2.FileDescriptorProto()
  103. file_proto.name = os.path.join(package.replace('.', '/'), proto_file_name)
  104. file_proto.package = package
  105. desc_proto = file_proto.message_type.add()
  106. desc_proto.name = name
  107. for f_number, (f_name, f_type) in enumerate(field_items, 1):
  108. field_proto = desc_proto.field.add()
  109. field_proto.name = f_name
  110. field_proto.number = f_number
  111. field_proto.label = descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL
  112. field_proto.type = f_type
  113. return file_proto