proto_builder.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117
  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. """Dynamic Protobuf class creator."""
  8. from collections import OrderedDict
  9. import hashlib
  10. import os
  11. from google.protobuf import descriptor
  12. from google.protobuf import descriptor_pb2
  13. from google.protobuf import descriptor_pool
  14. from google.protobuf import message_factory
  15. def _GetMessageFromFactory(pool, full_name):
  16. """Get a proto class from the MessageFactory by name.
  17. Args:
  18. pool: a descriptor pool.
  19. full_name: str, the fully qualified name of the proto type.
  20. Returns:
  21. A class, for the type identified by full_name.
  22. Raises:
  23. KeyError, if the proto is not found in the factory's descriptor pool.
  24. """
  25. proto_descriptor = pool.FindMessageTypeByName(full_name)
  26. proto_cls = message_factory.GetMessageClass(proto_descriptor)
  27. return proto_cls
  28. def MakeSimpleProtoClass(fields, full_name=None, pool=None):
  29. """Create a Protobuf class whose fields are basic types.
  30. Note: this doesn't validate field names!
  31. Args:
  32. fields: dict of {name: field_type} mappings for each field in the proto. If
  33. this is an OrderedDict the order will be maintained, otherwise the fields
  34. will be sorted by name.
  35. full_name: optional str, the fully-qualified name of the proto type.
  36. pool: optional DescriptorPool instance.
  37. Returns:
  38. a class, the new protobuf class with a FileDescriptor.
  39. """
  40. pool_instance = pool or descriptor_pool.DescriptorPool()
  41. if full_name is not None:
  42. try:
  43. proto_cls = _GetMessageFromFactory(pool_instance, full_name)
  44. return proto_cls
  45. except KeyError:
  46. # The factory's DescriptorPool doesn't know about this class yet.
  47. pass
  48. # Get a list of (name, field_type) tuples from the fields dict. If fields was
  49. # an OrderedDict we keep the order, but otherwise we sort the field to ensure
  50. # consistent ordering.
  51. field_items = fields.items()
  52. if not isinstance(fields, OrderedDict):
  53. field_items = sorted(field_items)
  54. # Use a consistent file name that is unlikely to conflict with any imported
  55. # proto files.
  56. fields_hash = hashlib.sha1(usedforsecurity=False)
  57. for f_name, f_type in field_items:
  58. fields_hash.update(f_name.encode('utf-8'))
  59. fields_hash.update(str(f_type).encode('utf-8'))
  60. proto_file_name = fields_hash.hexdigest() + '.proto'
  61. # If the proto is anonymous, use the same hash to name it.
  62. if full_name is None:
  63. full_name = (
  64. 'net.proto2.python.public.proto_builder.AnonymousProto_'
  65. + fields_hash.hexdigest()
  66. )
  67. try:
  68. proto_cls = _GetMessageFromFactory(pool_instance, full_name)
  69. return proto_cls
  70. except KeyError:
  71. # The factory's DescriptorPool doesn't know about this class yet.
  72. pass
  73. # This is the first time we see this proto: add a new descriptor to the pool.
  74. pool_instance.Add(
  75. _MakeFileDescriptorProto(proto_file_name, full_name, field_items)
  76. )
  77. return _GetMessageFromFactory(pool_instance, full_name)
  78. def _MakeFileDescriptorProto(proto_file_name, full_name, field_items):
  79. """Populate FileDescriptorProto for MessageFactory's DescriptorPool."""
  80. package, name = full_name.rsplit('.', 1)
  81. file_proto = descriptor_pb2.FileDescriptorProto()
  82. file_proto.name = os.path.join(package.replace('.', '/'), proto_file_name)
  83. file_proto.package = package
  84. desc_proto = file_proto.message_type.add()
  85. desc_proto.name = name
  86. for f_number, (f_name, f_type) in enumerate(field_items, 1):
  87. field_proto = desc_proto.field.add()
  88. field_proto.name = f_name
  89. # # If the number falls in the reserved range, reassign it to the correct
  90. # # number after the range.
  91. if f_number >= descriptor.FieldDescriptor.FIRST_RESERVED_FIELD_NUMBER:
  92. f_number += (
  93. descriptor.FieldDescriptor.LAST_RESERVED_FIELD_NUMBER
  94. - descriptor.FieldDescriptor.FIRST_RESERVED_FIELD_NUMBER
  95. + 1
  96. )
  97. field_proto.number = f_number
  98. field_proto.label = descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL
  99. field_proto.type = f_type
  100. return file_proto