descriptor_database.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185
  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. """Provides a container for DescriptorProtos."""
  8. __author__ = 'matthewtoia@google.com (Matt Toia)'
  9. from typing import Dict, Iterator, Optional
  10. import warnings
  11. class Error(Exception):
  12. pass
  13. class DescriptorDatabaseConflictingDefinitionError(Error):
  14. """Raised when a proto is added with the same name & different descriptor."""
  15. class DescriptorDatabase(object):
  16. """A container accepting FileDescriptorProtos and maps DescriptorProtos."""
  17. def __init__(self) -> None:
  18. self._file_desc_protos_by_file: Dict[
  19. str, 'descriptor_pb2.FileDescriptorProto'
  20. ] = {}
  21. self._file_desc_protos_by_symbol: Dict[
  22. str, 'descriptor_pb2.FileDescriptorProto'
  23. ] = {}
  24. def Add(self, file_desc_proto: 'descriptor_pb2.FileDescriptorProto') -> None:
  25. """Adds the FileDescriptorProto and its types to this database.
  26. Args:
  27. file_desc_proto: The FileDescriptorProto to add.
  28. Raises:
  29. DescriptorDatabaseConflictingDefinitionError: if an attempt is made to
  30. add a proto with the same name but different definition than an
  31. existing proto in the database.
  32. """
  33. proto_name = file_desc_proto.name
  34. if proto_name not in self._file_desc_protos_by_file:
  35. self._file_desc_protos_by_file[proto_name] = file_desc_proto
  36. elif self._file_desc_protos_by_file[proto_name] != file_desc_proto:
  37. raise DescriptorDatabaseConflictingDefinitionError(
  38. '%s already added, but with different descriptor.' % proto_name)
  39. else:
  40. return
  41. # Add all the top-level descriptors to the index.
  42. package = file_desc_proto.package
  43. for message in file_desc_proto.message_type:
  44. for name in _ExtractSymbols(message, package):
  45. self._AddSymbol(name, file_desc_proto)
  46. for enum in file_desc_proto.enum_type:
  47. self._AddSymbol(
  48. ('.'.join((package, enum.name)) if package else enum.name),
  49. file_desc_proto,
  50. )
  51. for enum_value in enum.value:
  52. self._file_desc_protos_by_symbol[
  53. '.'.join((package, enum_value.name)) if package else enum_value.name
  54. ] = file_desc_proto
  55. for extension in file_desc_proto.extension:
  56. self._AddSymbol(
  57. ('.'.join((package, extension.name)) if package else extension.name),
  58. file_desc_proto,
  59. )
  60. for service in file_desc_proto.service:
  61. self._AddSymbol(
  62. ('.'.join((package, service.name)) if package else service.name),
  63. file_desc_proto,
  64. )
  65. def FindFileByName(self, name: str) -> 'descriptor_pb2.FileDescriptorProto':
  66. """Finds the file descriptor proto by file name.
  67. Typically the file name is a relative path ending to a .proto file. The
  68. proto with the given name will have to have been added to this database
  69. using the Add method or else an error will be raised.
  70. Args:
  71. name: The file name to find.
  72. Returns:
  73. The file descriptor proto matching the name.
  74. Raises:
  75. KeyError if no file by the given name was added.
  76. """
  77. return self._file_desc_protos_by_file[name]
  78. def FindFileContainingSymbol(
  79. self, symbol: str
  80. ) -> 'descriptor_pb2.FileDescriptorProto':
  81. """Finds the file descriptor proto containing the specified symbol.
  82. The symbol should be a fully qualified name including the file descriptor's
  83. package and any containing messages. Some examples:
  84. 'some.package.name.Message'
  85. 'some.package.name.Message.NestedEnum'
  86. 'some.package.name.Message.some_field'
  87. The file descriptor proto containing the specified symbol must be added to
  88. this database using the Add method or else an error will be raised.
  89. Args:
  90. symbol: The fully qualified symbol name.
  91. Returns:
  92. The file descriptor proto containing the symbol.
  93. Raises:
  94. KeyError if no file contains the specified symbol.
  95. """
  96. if symbol.count('.') == 1 and symbol[0] == '.':
  97. symbol = symbol.lstrip('.')
  98. warnings.warn(
  99. 'Please remove the leading "." when '
  100. 'FindFileContainingSymbol, this will turn to error '
  101. 'in 2026 Jan.',
  102. RuntimeWarning,
  103. )
  104. try:
  105. return self._file_desc_protos_by_symbol[symbol]
  106. except KeyError:
  107. # Fields, enum values, and nested extensions are not in
  108. # _file_desc_protos_by_symbol. Try to find the top level
  109. # descriptor. Non-existent nested symbol under a valid top level
  110. # descriptor can also be found. The behavior is the same with
  111. # protobuf C++.
  112. top_level, _, _ = symbol.rpartition('.')
  113. try:
  114. return self._file_desc_protos_by_symbol[top_level]
  115. except KeyError:
  116. # Raise the original symbol as a KeyError for better diagnostics.
  117. raise KeyError(symbol)
  118. def FindFileContainingExtension(
  119. self, extendee_name: str, extension_number: int # pylint: disable=unused-argument
  120. ) -> Optional['descriptor_pb2.FileDescriptorProto']:
  121. # TODO: implement this API.
  122. return None
  123. def FindAllExtensionNumbers(self, extendee_name: str) -> list[int]: # pylint: disable=unused-argument
  124. # TODO: implement this API.
  125. return []
  126. def _AddSymbol(
  127. self, name: str, file_desc_proto: 'descriptor_pb2.FileDescriptorProto'
  128. ) -> None:
  129. if name in self._file_desc_protos_by_symbol:
  130. warn_msg = ('Conflict register for file "' + file_desc_proto.name +
  131. '": ' + name +
  132. ' is already defined in file "' +
  133. self._file_desc_protos_by_symbol[name].name + '"')
  134. warnings.warn(warn_msg, RuntimeWarning)
  135. self._file_desc_protos_by_symbol[name] = file_desc_proto
  136. def _ExtractSymbols(
  137. desc_proto: 'descriptor_pb2.DescriptorProto', package: str
  138. ) -> Iterator[str]:
  139. """Pulls out all the symbols from a descriptor proto.
  140. Args:
  141. desc_proto: The proto to extract symbols from.
  142. package: The package containing the descriptor type.
  143. Yields:
  144. The fully qualified name found in the descriptor.
  145. """
  146. message_name = package + '.' + desc_proto.name if package else desc_proto.name
  147. yield message_name
  148. for nested_type in desc_proto.nested_type:
  149. for symbol in _ExtractSymbols(nested_type, message_name):
  150. yield symbol
  151. for enum_type in desc_proto.enum_type:
  152. yield '.'.join((message_name, enum_type.name))