symbol_database.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197
  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 GetPrototype(self, descriptor):
  40. warnings.warn('SymbolDatabase.GetPrototype() is deprecated. Please '
  41. 'use message_factory.GetMessageClass() instead. '
  42. 'SymbolDatabase.GetPrototype() will be removed soon.')
  43. return message_factory.GetMessageClass(descriptor)
  44. def CreatePrototype(self, descriptor):
  45. warnings.warn('Directly call CreatePrototype() is wrong. Please use '
  46. 'message_factory.GetMessageClass() instead. '
  47. 'SymbolDatabase.CreatePrototype() will be removed soon.')
  48. return message_factory._InternalCreateMessageClass(descriptor)
  49. def GetMessages(self, files):
  50. warnings.warn('SymbolDatabase.GetMessages() is deprecated. Please use '
  51. 'message_factory.GetMessageClassedForFiles() instead. '
  52. 'SymbolDatabase.GetMessages() will be removed soon.')
  53. return message_factory.GetMessageClassedForFiles(files, self.pool)
  54. def RegisterMessage(self, message):
  55. """Registers the given message type in the local database.
  56. Calls to GetSymbol() and GetMessages() will return messages registered here.
  57. Args:
  58. message: A :class:`google.protobuf.message.Message` subclass (or
  59. instance); its descriptor will be registered.
  60. Returns:
  61. The provided message.
  62. """
  63. desc = message.DESCRIPTOR
  64. self._classes[desc] = message
  65. self.RegisterMessageDescriptor(desc)
  66. return message
  67. def RegisterMessageDescriptor(self, message_descriptor):
  68. """Registers the given message descriptor in the local database.
  69. Args:
  70. message_descriptor (Descriptor): the message descriptor to add.
  71. """
  72. if api_implementation.Type() == 'python':
  73. # pylint: disable=protected-access
  74. self.pool._AddDescriptor(message_descriptor)
  75. def RegisterEnumDescriptor(self, enum_descriptor):
  76. """Registers the given enum descriptor in the local database.
  77. Args:
  78. enum_descriptor (EnumDescriptor): The enum descriptor to register.
  79. Returns:
  80. EnumDescriptor: The provided descriptor.
  81. """
  82. if api_implementation.Type() == 'python':
  83. # pylint: disable=protected-access
  84. self.pool._AddEnumDescriptor(enum_descriptor)
  85. return enum_descriptor
  86. def RegisterServiceDescriptor(self, service_descriptor):
  87. """Registers the given service descriptor in the local database.
  88. Args:
  89. service_descriptor (ServiceDescriptor): the service descriptor to
  90. register.
  91. """
  92. if api_implementation.Type() == 'python':
  93. # pylint: disable=protected-access
  94. self.pool._AddServiceDescriptor(service_descriptor)
  95. def RegisterFileDescriptor(self, file_descriptor):
  96. """Registers the given file descriptor in the local database.
  97. Args:
  98. file_descriptor (FileDescriptor): The file descriptor to register.
  99. """
  100. if api_implementation.Type() == 'python':
  101. # pylint: disable=protected-access
  102. self.pool._InternalAddFileDescriptor(file_descriptor)
  103. def GetSymbol(self, symbol):
  104. """Tries to find a symbol in the local database.
  105. Currently, this method only returns message.Message instances, however, if
  106. may be extended in future to support other symbol types.
  107. Args:
  108. symbol (str): a protocol buffer symbol.
  109. Returns:
  110. A Python class corresponding to the symbol.
  111. Raises:
  112. KeyError: if the symbol could not be found.
  113. """
  114. return self._classes[self.pool.FindMessageTypeByName(symbol)]
  115. def GetMessages(self, files):
  116. # TODO: Fix the differences with MessageFactory.
  117. """Gets all registered messages from a specified file.
  118. Only messages already created and registered will be returned; (this is the
  119. case for imported _pb2 modules)
  120. But unlike MessageFactory, this version also returns already defined nested
  121. messages, but does not register any message extensions.
  122. Args:
  123. files (list[str]): The file names to extract messages from.
  124. Returns:
  125. A dictionary mapping proto names to the message classes.
  126. Raises:
  127. KeyError: if a file could not be found.
  128. """
  129. def _GetAllMessages(desc):
  130. """Walk a message Descriptor and recursively yields all message names."""
  131. yield desc
  132. for msg_desc in desc.nested_types:
  133. for nested_desc in _GetAllMessages(msg_desc):
  134. yield nested_desc
  135. result = {}
  136. for file_name in files:
  137. file_desc = self.pool.FindFileByName(file_name)
  138. for msg_desc in file_desc.message_types_by_name.values():
  139. for desc in _GetAllMessages(msg_desc):
  140. try:
  141. result[desc.full_name] = self._classes[desc]
  142. except KeyError:
  143. # This descriptor has no registered class, skip it.
  144. pass
  145. return result
  146. _DEFAULT = SymbolDatabase(pool=descriptor_pool.Default())
  147. def Default():
  148. """Returns the default SymbolDatabase."""
  149. return _DEFAULT