symbol_database.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179
  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. """A database of Python protocol buffer generated symbols.
  8. SymbolDatabase is the MessageFactory for messages generated at compile time,
  9. and makes it easy to create new instances of a registered type, given only the
  10. type's protocol buffer symbol name.
  11. Example usage::
  12. db = symbol_database.SymbolDatabase()
  13. # Register symbols of interest, from one or multiple files.
  14. db.RegisterFileDescriptor(my_proto_pb2.DESCRIPTOR)
  15. db.RegisterMessage(my_proto_pb2.MyMessage)
  16. db.RegisterEnumDescriptor(my_proto_pb2.MyEnum.DESCRIPTOR)
  17. # The database can be used as a MessageFactory, to generate types based on
  18. # their name:
  19. types = db.GetMessages(['my_proto.proto'])
  20. my_message_instance = types['MyMessage']()
  21. # The database's underlying descriptor pool can be queried, so it's not
  22. # necessary to know a type's filename to be able to generate it:
  23. filename = db.pool.FindFileContainingSymbol('MyMessage')
  24. my_message_instance = db.GetMessages([filename])['MyMessage']()
  25. # This functionality is also provided directly via a convenience method:
  26. my_message_instance = db.GetSymbol('MyMessage')()
  27. """
  28. import warnings
  29. from google.protobuf.internal import api_implementation
  30. from google.protobuf import descriptor_pool
  31. from google.protobuf import message_factory
  32. class SymbolDatabase():
  33. """A database of Python generated symbols."""
  34. # local cache of registered classes.
  35. _classes = {}
  36. def __init__(self, pool=None):
  37. """Initializes a new SymbolDatabase."""
  38. self.pool = pool or descriptor_pool.DescriptorPool()
  39. def RegisterMessage(self, message):
  40. """Registers the given message type in the local database.
  41. Calls to GetSymbol() and GetMessages() will return messages registered here.
  42. Args:
  43. message: A :class:`google.protobuf.message.Message` subclass (or
  44. instance); its descriptor will be registered.
  45. Returns:
  46. The provided message.
  47. """
  48. desc = message.DESCRIPTOR
  49. self._classes[desc] = message
  50. self.RegisterMessageDescriptor(desc)
  51. return message
  52. def RegisterMessageDescriptor(self, message_descriptor):
  53. """Registers the given message descriptor in the local database.
  54. Args:
  55. message_descriptor (Descriptor): the message descriptor to add.
  56. """
  57. if api_implementation.Type() == 'python':
  58. # pylint: disable=protected-access
  59. self.pool._AddDescriptor(message_descriptor)
  60. def RegisterEnumDescriptor(self, enum_descriptor):
  61. """Registers the given enum descriptor in the local database.
  62. Args:
  63. enum_descriptor (EnumDescriptor): The enum descriptor to register.
  64. Returns:
  65. EnumDescriptor: The provided descriptor.
  66. """
  67. if api_implementation.Type() == 'python':
  68. # pylint: disable=protected-access
  69. self.pool._AddEnumDescriptor(enum_descriptor)
  70. return enum_descriptor
  71. def RegisterServiceDescriptor(self, service_descriptor):
  72. """Registers the given service descriptor in the local database.
  73. Args:
  74. service_descriptor (ServiceDescriptor): the service descriptor to
  75. register.
  76. """
  77. if api_implementation.Type() == 'python':
  78. # pylint: disable=protected-access
  79. self.pool._AddServiceDescriptor(service_descriptor)
  80. def RegisterFileDescriptor(self, file_descriptor):
  81. """Registers the given file descriptor in the local database.
  82. Args:
  83. file_descriptor (FileDescriptor): The file descriptor to register.
  84. """
  85. if api_implementation.Type() == 'python':
  86. # pylint: disable=protected-access
  87. self.pool._InternalAddFileDescriptor(file_descriptor)
  88. def GetSymbol(self, symbol):
  89. """Tries to find a symbol in the local database.
  90. Currently, this method only returns message.Message instances, however, if
  91. may be extended in future to support other symbol types.
  92. Args:
  93. symbol (str): a protocol buffer symbol.
  94. Returns:
  95. A Python class corresponding to the symbol.
  96. Raises:
  97. KeyError: if the symbol could not be found.
  98. """
  99. return self._classes[self.pool.FindMessageTypeByName(symbol)]
  100. def GetMessages(self, files):
  101. # TODO: Fix the differences with MessageFactory.
  102. """Gets all registered messages from a specified file.
  103. Only messages already created and registered will be returned; (this is the
  104. case for imported _pb2 modules)
  105. But unlike MessageFactory, this version also returns already defined nested
  106. messages, but does not register any message extensions.
  107. Args:
  108. files (list[str]): The file names to extract messages from.
  109. Returns:
  110. A dictionary mapping proto names to the message classes.
  111. Raises:
  112. KeyError: if a file could not be found.
  113. """
  114. def _GetAllMessages(desc):
  115. """Walk a message Descriptor and recursively yields all message names."""
  116. yield desc
  117. for msg_desc in desc.nested_types:
  118. for nested_desc in _GetAllMessages(msg_desc):
  119. yield nested_desc
  120. result = {}
  121. for file_name in files:
  122. file_desc = self.pool.FindFileByName(file_name)
  123. for msg_desc in file_desc.message_types_by_name.values():
  124. for desc in _GetAllMessages(msg_desc):
  125. try:
  126. result[desc.full_name] = self._classes[desc]
  127. except KeyError:
  128. # This descriptor has no registered class, skip it.
  129. pass
  130. return result
  131. _DEFAULT = SymbolDatabase(pool=descriptor_pool.Default())
  132. def Default():
  133. """Returns the default SymbolDatabase."""
  134. return _DEFAULT