extension_dict.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209
  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. """Contains _ExtensionDict class to represent extensions."""
  8. from google.protobuf.descriptor import FieldDescriptor
  9. from google.protobuf.internal import type_checkers
  10. def _VerifyExtensionHandle(message, extension_handle):
  11. """Verify that the given extension handle is valid."""
  12. if not isinstance(extension_handle, FieldDescriptor):
  13. raise KeyError(
  14. 'HasExtension() expects an extension handle, got: %s' % extension_handle
  15. )
  16. if not extension_handle.is_extension:
  17. raise KeyError('"%s" is not an extension.' % extension_handle.full_name)
  18. if not extension_handle.containing_type:
  19. raise KeyError(
  20. '"%s" is missing a containing_type.' % extension_handle.full_name
  21. )
  22. if extension_handle.containing_type is not message.DESCRIPTOR:
  23. raise KeyError(
  24. 'Extension "%s" extends message type "%s", but this '
  25. 'message is of type "%s".'
  26. % (
  27. extension_handle.full_name,
  28. extension_handle.containing_type.full_name,
  29. message.DESCRIPTOR.full_name,
  30. )
  31. )
  32. # TODO: Unify error handling of "unknown extension" crap.
  33. # TODO: Support iteritems()-style iteration over all
  34. # extensions with the "has" bits turned on?
  35. class _ExtensionDict(object):
  36. """Dict-like container for Extension fields on proto instances.
  37. Note that in all cases we expect extension handles to be
  38. FieldDescriptors.
  39. """
  40. def __init__(self, extended_message):
  41. """Args:
  42. extended_message: Message instance for which we are the Extensions dict.
  43. """
  44. self._extended_message = extended_message
  45. def __getitem__(self, extension_handle):
  46. """Returns the current value of the given extension handle."""
  47. _VerifyExtensionHandle(self._extended_message, extension_handle)
  48. result = self._extended_message._fields.get(extension_handle)
  49. if result is not None:
  50. return result
  51. if extension_handle.is_repeated:
  52. result = extension_handle._default_constructor(self._extended_message)
  53. elif extension_handle.cpp_type == FieldDescriptor.CPPTYPE_MESSAGE:
  54. message_type = extension_handle.message_type
  55. if not hasattr(message_type, '_concrete_class'):
  56. # pylint: disable=g-import-not-at-top
  57. from google.protobuf import message_factory
  58. message_factory.GetMessageClass(message_type)
  59. if not hasattr(extension_handle.message_type, '_concrete_class'):
  60. from google.protobuf import message_factory
  61. message_factory.GetMessageClass(extension_handle.message_type)
  62. result = extension_handle.message_type._concrete_class()
  63. try:
  64. result._SetListener(self._extended_message._listener_for_children)
  65. except ReferenceError:
  66. pass
  67. else:
  68. # Singular scalar -- just return the default without inserting into the
  69. # dict.
  70. return extension_handle.default_value
  71. # Atomically check if another thread has preempted us and, if not, swap
  72. # in the new object we just created. If someone has preempted us, we
  73. # take that object and discard ours.
  74. # WARNING: We are relying on setdefault() being atomic. This is true
  75. # in CPython but we haven't investigated others. This warning appears
  76. # in several other locations in this file.
  77. if self._extended_message._frozen:
  78. result._SetFrozen()
  79. result = self._extended_message._fields.setdefault(extension_handle, result)
  80. return result
  81. def __eq__(self, other):
  82. if not isinstance(other, self.__class__):
  83. return False
  84. my_fields = self._extended_message.ListFields()
  85. other_fields = other._extended_message.ListFields()
  86. # Get rid of non-extension fields.
  87. my_fields = [field for field in my_fields if field.is_extension]
  88. other_fields = [field for field in other_fields if field.is_extension]
  89. return my_fields == other_fields
  90. def __ne__(self, other):
  91. return not self == other
  92. def __len__(self):
  93. fields = self._extended_message.ListFields()
  94. # Get rid of non-extension fields.
  95. extension_fields = [field for field in fields if field[0].is_extension]
  96. return len(extension_fields)
  97. def __hash__(self):
  98. raise TypeError('unhashable object')
  99. # Note that this is only meaningful for non-repeated, scalar extension
  100. # fields. Note also that we may have to call _Modified() when we do
  101. # successfully set a field this way, to set any necessary "has" bits in the
  102. # ancestors of the extended message.
  103. def __setitem__(self, extension_handle, value):
  104. """If extension_handle specifies a non-repeated, scalar extension
  105. field, sets the value of that field.
  106. """
  107. _VerifyExtensionHandle(self._extended_message, extension_handle)
  108. self._extended_message._AssureWritable()
  109. if (
  110. extension_handle.is_repeated
  111. or extension_handle.cpp_type == FieldDescriptor.CPPTYPE_MESSAGE
  112. ):
  113. raise TypeError(
  114. 'Cannot assign to extension "%s" because it is a repeated or '
  115. 'composite type.'
  116. % extension_handle.full_name
  117. )
  118. # It's slightly wasteful to lookup the type checker each time,
  119. # but we expect this to be a vanishingly uncommon case anyway.
  120. type_checker = type_checkers.GetTypeChecker(extension_handle)
  121. # pylint: disable=protected-access
  122. self._extended_message._fields[extension_handle] = type_checker.CheckValue(
  123. value
  124. )
  125. self._extended_message._Modified()
  126. def __delitem__(self, extension_handle):
  127. self._extended_message.ClearExtension(extension_handle)
  128. def _FindExtensionByName(self, name):
  129. """Tries to find a known extension with the specified name.
  130. Args:
  131. name: Extension full name.
  132. Returns:
  133. Extension field descriptor.
  134. """
  135. descriptor = self._extended_message.DESCRIPTOR
  136. extensions = descriptor.file.pool._extensions_by_name[descriptor]
  137. return extensions.get(name, None)
  138. def _FindExtensionByNumber(self, number):
  139. """Tries to find a known extension with the field number.
  140. Args:
  141. number: Extension field number.
  142. Returns:
  143. Extension field descriptor.
  144. """
  145. descriptor = self._extended_message.DESCRIPTOR
  146. extensions = descriptor.file.pool._extensions_by_number[descriptor]
  147. return extensions.get(number, None)
  148. def __iter__(self):
  149. # Return a generator over the populated extension fields
  150. return (
  151. f[0] for f in self._extended_message.ListFields() if f[0].is_extension
  152. )
  153. def __contains__(self, extension_handle):
  154. _VerifyExtensionHandle(self._extended_message, extension_handle)
  155. if extension_handle not in self._extended_message._fields:
  156. return False
  157. if extension_handle.is_repeated:
  158. return bool(self._extended_message._fields.get(extension_handle))
  159. if extension_handle.cpp_type == FieldDescriptor.CPPTYPE_MESSAGE:
  160. value = self._extended_message._fields.get(extension_handle)
  161. # pylint: disable=protected-access
  162. return value is not None and value._is_present_in_parent
  163. return True