third_party/protobuf: Update to HEAD (83d681ee2c)

third_party/protobuf/python/google/protobuf/descriptor_pool.py

 # Protocol Buffers - Google's data interchange format
 # Copyright 2008 Google Inc.  All rights reserved.
 # https://developers.google.com/protocol-buffers/
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are
 # met:
 #
 #     * Redistributions of source code must retain the above copyright
 # notice, this list of conditions and the following disclaimer.
@@ skipping 39 matching lines @@
 module in order to create a protocol buffer class that can be encoded and
 decoded.
 
 If you want to get a Python class for the specified proto, use the
 helper functions inside google.protobuf.message_factory
 directly instead of this class.
 """
 
 __author__ = 'matthewtoia@google.com (Matt Toia)'
 
+import collections
+
 from google.protobuf import descriptor
 from google.protobuf import descriptor_database
 from google.protobuf import text_encoding
 
 
-_USE_C_DESCRIPTORS = descriptor._USE_C_DESCRIPTORS
+_USE_C_DESCRIPTORS = descriptor._USE_C_DESCRIPTORS  # pylint: disable=protected-access
 
 
 def _NormalizeFullyQualifiedName(name):
   """Remove leading period from fully-qualified type name.
 
   Due to b/13860351 in descriptor_database.py, types in the root namespace are
   generated with a leading period. This function removes that prefix.
 
   Args:
     name: A str, the fully-qualified symbol name.
 
   Returns:
     A str, the normalized fully-qualified symbol name.
   """
   return name.lstrip('.')
 
 
+def _OptionsOrNone(descriptor_proto):
+  """Returns the value of the field `options`, or None if it is not set."""
+  if descriptor_proto.HasField('options'):
+    return descriptor_proto.options
+  else:
+    return None
+
+
+def _IsMessageSetExtension(field):
+  return (field.is_extension and
+          field.containing_type.has_options and
+          field.containing_type.GetOptions().message_set_wire_format and
+          field.type == descriptor.FieldDescriptor.TYPE_MESSAGE and
+          field.label == descriptor.FieldDescriptor.LABEL_OPTIONAL)
+
+
 class DescriptorPool(object):
   """A collection of protobufs dynamically constructed by descriptor protos."""
 
   if _USE_C_DESCRIPTORS:
 
     def __new__(cls, descriptor_db=None):
       # pylint: disable=protected-access
       return descriptor._message.DescriptorPool(descriptor_db)
 
   def __init__(self, descriptor_db=None):
     """Initializes a Pool of proto buffs.
 
     The descriptor_db argument to the constructor is provided to allow
     specialized file descriptor proto lookup code to be triggered on demand. An
     example would be an implementation which will read and compile a file
     specified in a call to FindFileByName() and not require the call to Add()
     at all. Results from this database will be cached internally here as well.
 
     Args:
       descriptor_db: A secondary source of file descriptors.
     """
 
     self._internal_db = descriptor_database.DescriptorDatabase()
     self._descriptor_db = descriptor_db
     self._descriptors = {}
     self._enum_descriptors = {}
     self._file_descriptors = {}
+    self._toplevel_extensions = {}
+    # We store extensions in two two-level mappings: The first key is the
+    # descriptor of the message being extended, the second key is the extension
+    # full name or its tag number.
+    self._extensions_by_name = collections.defaultdict(dict)
+    self._extensions_by_number = collections.defaultdict(dict)
 
   def Add(self, file_desc_proto):
     """Adds the FileDescriptorProto and its types to this pool.
 
     Args:
       file_desc_proto: The FileDescriptorProto to add.
     """
 
     self._internal_db.Add(file_desc_proto)
 
@@ skipping 35 matching lines @@
     Args:
       enum_desc: An EnumDescriptor.
     """
 
     if not isinstance(enum_desc, descriptor.EnumDescriptor):
       raise TypeError('Expected instance of descriptor.EnumDescriptor.')
 
     self._enum_descriptors[enum_desc.full_name] = enum_desc
     self.AddFileDescriptor(enum_desc.file)
 
+  def AddExtensionDescriptor(self, extension):
+    """Adds a FieldDescriptor describing an extension to the pool.
+
+    Args:
+      extension: A FieldDescriptor.
+
+    Raises:
+      AssertionError: when another extension with the same number extends the
+        same message.
+      TypeError: when the specified extension is not a
+        descriptor.FieldDescriptor.
+    """
+    if not (isinstance(extension, descriptor.FieldDescriptor) and
+            extension.is_extension):
+      raise TypeError('Expected an extension descriptor.')
+
+    if extension.extension_scope is None:
+      self._toplevel_extensions[extension.full_name] = extension
+
+    try:
+      existing_desc = self._extensions_by_number[
+          extension.containing_type][extension.number]
+    except KeyError:
+      pass
+    else:
+      if extension is not existing_desc:
+        raise AssertionError(
+            'Extensions "%s" and "%s" both try to extend message type "%s" '
+            'with field number %d.' %
+            (extension.full_name, existing_desc.full_name,
+             extension.containing_type.full_name, extension.number))
+
+    self._extensions_by_number[extension.containing_type][
+        extension.number] = extension
+    self._extensions_by_name[extension.containing_type][
+        extension.full_name] = extension
+
+    # Also register MessageSet extensions with the type name.
+    if _IsMessageSetExtension(extension):
+      self._extensions_by_name[extension.containing_type][
+          extension.message_type.full_name] = extension
+
   def AddFileDescriptor(self, file_desc):
     """Adds a FileDescriptor to the pool, non-recursively.
 
     If the FileDescriptor contains messages or enums, the caller must explicitly
     register them.
 
     Args:
       file_desc: A FileDescriptor.
     """
 
@@ skipping 112 matching lines @@
   def FindExtensionByName(self, full_name):
     """Loads the named extension descriptor from the pool.
 
     Args:
       full_name: The full name of the extension descriptor to load.
 
     Returns:
       A FieldDescriptor, describing the named extension.
     """
     full_name = _NormalizeFullyQualifiedName(full_name)
+    try:
+      # The proto compiler does not give any link between the FileDescriptor
+      # and top-level extensions unless the FileDescriptorProto is added to
+      # the DescriptorDatabase, but this can impact memory usage.
+      # So we registered these extensions by name explicitly.
+      return self._toplevel_extensions[full_name]
+    except KeyError:
+      pass
     message_name, _, extension_name = full_name.rpartition('.')
     try:
       # Most extensions are nested inside a message.
       scope = self.FindMessageTypeByName(message_name)
     except KeyError:
       # Some extensions are defined at file scope.
       scope = self.FindFileContainingSymbol(full_name)
     return scope.extensions_by_name[extension_name]
 
+  def FindExtensionByNumber(self, message_descriptor, number):
+    """Gets the extension of the specified message with the specified number.
+
+    Extensions have to be registered to this pool by calling
+    AddExtensionDescriptor.
+
+    Args:
+      message_descriptor: descriptor of the extended message.
+      number: integer, number of the extension field.
+
+    Returns:
+      A FieldDescriptor describing the extension.
+
+    Raise:
+      KeyError: when no extension with the given number is known for the
+        specified message.
+    """
+    return self._extensions_by_number[message_descriptor][number]
+
+  def FindAllExtensions(self, message_descriptor):
+    """Gets all the known extension of a given message.
+
+    Extensions have to be registered to this pool by calling
+    AddExtensionDescriptor.
+
+    Args:
+      message_descriptor: descriptor of the extended message.
+
+    Returns:
+      A list of FieldDescriptor describing the extensions.
+    """
+    return list(self._extensions_by_number[message_descriptor].values())
+
   def _ConvertFileProtoToFileDescriptor(self, file_proto):
     """Creates a FileDescriptor from a proto or returns a cached copy.
 
     This method also has the side effect of loading all the symbols found in
     the file into the appropriate dictionaries in the pool.
 
     Args:
       file_proto: The proto to convert.
 
     Returns:
       A FileDescriptor matching the passed in proto.
     """
 
     if file_proto.name not in self._file_descriptors:
       built_deps = list(self._GetDeps(file_proto.dependency))
       direct_deps = [self.FindFileByName(n) for n in file_proto.dependency]
       public_deps = [direct_deps[i] for i in file_proto.public_dependency]
 
       file_descriptor = descriptor.FileDescriptor(
           pool=self,
           name=file_proto.name,
           package=file_proto.package,
           syntax=file_proto.syntax,
-          options=file_proto.options,
+          options=_OptionsOrNone(file_proto),
           serialized_pb=file_proto.SerializeToString(),
           dependencies=direct_deps,
           public_dependencies=public_deps)
-      if _USE_C_DESCRIPTORS:
-        # When using C++ descriptors, all objects defined in the file were added
-        # to the C++ database when the FileDescriptor was built above.
-        # Just add them to this descriptor pool.
-        def _AddMessageDescriptor(message_desc):
-          self._descriptors[message_desc.full_name] = message_desc
-          for nested in message_desc.nested_types:
-            _AddMessageDescriptor(nested)
-          for enum_type in message_desc.enum_types:
-            _AddEnumDescriptor(enum_type)
-        def _AddEnumDescriptor(enum_desc):
-          self._enum_descriptors[enum_desc.full_name] = enum_desc
-        for message_type in file_descriptor.message_types_by_name.values():
-          _AddMessageDescriptor(message_type)
-        for enum_type in file_descriptor.enum_types_by_name.values():
-          _AddEnumDescriptor(enum_type)
+      scope = {}
+
+      # This loop extracts all the message and enum types from all the
+      # dependencies of the file_proto. This is necessary to create the
+      # scope of available message types when defining the passed in
+      # file proto.
+      for dependency in built_deps:
+        scope.update(self._ExtractSymbols(
+            dependency.message_types_by_name.values()))
+        scope.update((_PrefixWithDot(enum.full_name), enum)
+                     for enum in dependency.enum_types_by_name.values())
+
+      for message_type in file_proto.message_type:
+        message_desc = self._ConvertMessageDescriptor(
+            message_type, file_proto.package, file_descriptor, scope,
+            file_proto.syntax)
+        file_descriptor.message_types_by_name[message_desc.name] = (
+            message_desc)
+
+      for enum_type in file_proto.enum_type:
+        file_descriptor.enum_types_by_name[enum_type.name] = (
+            self._ConvertEnumDescriptor(enum_type, file_proto.package,
+                                        file_descriptor, None, scope))
+
+      for index, extension_proto in enumerate(file_proto.extension):
+        extension_desc = self._MakeFieldDescriptor(
+            extension_proto, file_proto.package, index, is_extension=True)
+        extension_desc.containing_type = self._GetTypeFromScope(
+            file_descriptor.package, extension_proto.extendee, scope)
+        self._SetFieldType(extension_proto, extension_desc,
+                           file_descriptor.package, scope)
+        file_descriptor.extensions_by_name[extension_desc.name] = (
+            extension_desc)
+
+      for desc_proto in file_proto.message_type:
+        self._SetAllFieldTypes(file_proto.package, desc_proto, scope)
+
+      if file_proto.package:
+        desc_proto_prefix = _PrefixWithDot(file_proto.package)
       else:
-        scope = {}
+        desc_proto_prefix = ''
 
-        # This loop extracts all the message and enum types from all the
-        # dependencies of the file_proto. This is necessary to create the
-        # scope of available message types when defining the passed in
-        # file proto.
-        for dependency in built_deps:
-          scope.update(self._ExtractSymbols(
-              dependency.message_types_by_name.values()))
-          scope.update((_PrefixWithDot(enum.full_name), enum)
-                       for enum in dependency.enum_types_by_name.values())
+      for desc_proto in file_proto.message_type:
+        desc = self._GetTypeFromScope(
+            desc_proto_prefix, desc_proto.name, scope)
+        file_descriptor.message_types_by_name[desc_proto.name] = desc
 
-        for message_type in file_proto.message_type:
-          message_desc = self._ConvertMessageDescriptor(
-              message_type, file_proto.package, file_descriptor, scope,
-              file_proto.syntax)
-          file_descriptor.message_types_by_name[message_desc.name] = (
-              message_desc)
-
-        for enum_type in file_proto.enum_type:
-          file_descriptor.enum_types_by_name[enum_type.name] = (
-              self._ConvertEnumDescriptor(enum_type, file_proto.package,
-                                          file_descriptor, None, scope))
-
-        for index, extension_proto in enumerate(file_proto.extension):
-          extension_desc = self._MakeFieldDescriptor(
-              extension_proto, file_proto.package, index, is_extension=True)
-          extension_desc.containing_type = self._GetTypeFromScope(
-              file_descriptor.package, extension_proto.extendee, scope)
-          self._SetFieldType(extension_proto, extension_desc,
-                            file_descriptor.package, scope)
-          file_descriptor.extensions_by_name[extension_desc.name] = (
-              extension_desc)
-
-        for desc_proto in file_proto.message_type:
-          self._SetAllFieldTypes(file_proto.package, desc_proto, scope)
-
-        if file_proto.package:
-          desc_proto_prefix = _PrefixWithDot(file_proto.package)
-        else:
-          desc_proto_prefix = ''
-
-        for desc_proto in file_proto.message_type:
-          desc = self._GetTypeFromScope(
-              desc_proto_prefix, desc_proto.name, scope)
-          file_descriptor.message_types_by_name[desc_proto.name] = desc
+      for index, service_proto in enumerate(file_proto.service):
+        file_descriptor.services_by_name[service_proto.name] = (
+            self._MakeServiceDescriptor(service_proto, index, scope,
+                                        file_proto.package, file_descriptor))
 
       self.Add(file_proto)
       self._file_descriptors[file_proto.name] = file_descriptor
 
     return self._file_descriptors[file_proto.name]
 
   def _ConvertMessageDescriptor(self, desc_proto, package=None, file_desc=None,
                                 scope=None, syntax=None):
     """Adds the proto to the pool in the specified package.
 
     Args:
       desc_proto: The descriptor_pb2.DescriptorProto protobuf message.
       package: The package the proto should be located in.
       file_desc: The file containing this message.
       scope: Dict mapping short and full symbols to message and enum types.
+      syntax: string indicating syntax of the file ("proto2" or "proto3")
 
     Returns:
       The added descriptor.
     """
 
     if package:
       desc_name = '.'.join((package, desc_proto.name))
     else:
       desc_name = desc_proto.name
 
@@ skipping 13 matching lines @@
         self._ConvertEnumDescriptor(enum, desc_name, file_desc, None, scope)
         for enum in desc_proto.enum_type]
     fields = [self._MakeFieldDescriptor(field, desc_name, index)
               for index, field in enumerate(desc_proto.field)]
     extensions = [
         self._MakeFieldDescriptor(extension, desc_name, index,
                                   is_extension=True)
         for index, extension in enumerate(desc_proto.extension)]
     oneofs = [
         descriptor.OneofDescriptor(desc.name, '.'.join((desc_name, desc.name)),
-                                   index, None, [])
+                                   index, None, [], desc.options)
         for index, desc in enumerate(desc_proto.oneof_decl)]
     extension_ranges = [(r.start, r.end) for r in desc_proto.extension_range]
     if extension_ranges:
       is_extendable = True
     else:
       is_extendable = False
     desc = descriptor.Descriptor(
         name=desc_proto.name,
         full_name=desc_name,
         filename=file_name,
         containing_type=None,
         fields=fields,
         oneofs=oneofs,
         nested_types=nested,
         enum_types=enums,
         extensions=extensions,
-        options=desc_proto.options,
+        options=_OptionsOrNone(desc_proto),
         is_extendable=is_extendable,
         extension_ranges=extension_ranges,
         file=file_desc,
         serialized_start=None,
         serialized_end=None,
         syntax=syntax)
     for nested in desc.nested_types:
       nested.containing_type = desc
     for enum in desc.enum_types:
       enum.containing_type = desc
@@ skipping 33 matching lines @@
       file_name = file_desc.name
 
     values = [self._MakeEnumValueDescriptor(value, index)
               for index, value in enumerate(enum_proto.value)]
     desc = descriptor.EnumDescriptor(name=enum_proto.name,
                                      full_name=enum_name,
                                      filename=file_name,
                                      file=file_desc,
                                      values=values,
                                      containing_type=containing_type,
-                                     options=enum_proto.options)
+                                     options=_OptionsOrNone(enum_proto))
     scope['.%s' % enum_name] = desc
     self._enum_descriptors[enum_name] = desc
     return desc
 
   def _MakeFieldDescriptor(self, field_proto, message_name, index,
                            is_extension=False):
     """Creates a field descriptor from a FieldDescriptorProto.
 
     For message and enum type fields, this method will do a look up
     in the pool for the appropriate descriptor for that type. If it
@@ skipping 24 matching lines @@
         type=field_proto.type,
         cpp_type=None,
         message_type=None,
         enum_type=None,
         containing_type=None,
         label=field_proto.label,
         has_default_value=False,
         default_value=None,
         is_extension=is_extension,
         extension_scope=None,
-        options=field_proto.options)
+        options=_OptionsOrNone(field_proto))
 
   def _SetAllFieldTypes(self, package, desc_proto, scope):
     """Sets all the descriptor's fields's types.
 
     This method also sets the containing types on any extensions.
 
     Args:
       package: The current package of desc_proto.
       desc_proto: The message descriptor to update.
       scope: Enclosing scope of available types.
@@ skipping 98 matching lines @@
       index: The index of the enum value.
 
     Returns:
       An initialized EnumValueDescriptor object.
     """
 
     return descriptor.EnumValueDescriptor(
         name=value_proto.name,
         index=index,
         number=value_proto.number,
-        options=value_proto.options,
+        options=_OptionsOrNone(value_proto),
         type=None)
 
+  def _MakeServiceDescriptor(self, service_proto, service_index, scope,
+                             package, file_desc):
+    """Make a protobuf ServiceDescriptor given a ServiceDescriptorProto.
+
+    Args:
+      service_proto: The descriptor_pb2.ServiceDescriptorProto protobuf message.
+      service_index: The index of the service in the File.
+      scope: Dict mapping short and full symbols to message and enum types.
+      package: Optional package name for the new message EnumDescriptor.
+      file_desc: The file containing the service descriptor.
+
+    Returns:
+      The added descriptor.
+    """
+
+    if package:
+      service_name = '.'.join((package, service_proto.name))
+    else:
+      service_name = service_proto.name
+
+    methods = [self._MakeMethodDescriptor(method_proto, service_name, package,
+                                          scope, index)
+               for index, method_proto in enumerate(service_proto.method)]
+    desc = descriptor.ServiceDescriptor(name=service_proto.name,
+                                        full_name=service_name,
+                                        index=service_index,
+                                        methods=methods,
+                                        options=_OptionsOrNone(service_proto),
+                                        file=file_desc)
+    return desc
+
+  def _MakeMethodDescriptor(self, method_proto, service_name, package, scope,
+                            index):
+    """Creates a method descriptor from a MethodDescriptorProto.
+
+    Args:
+      method_proto: The proto describing the method.
+      service_name: The name of the containing service.
+      package: Optional package name to look up for types.
+      scope: Scope containing available types.
+      index: Index of the method in the service.
+
+    Returns:
+      An initialized MethodDescriptor object.
+    """
+    full_name = '.'.join((service_name, method_proto.name))
+    input_type = self._GetTypeFromScope(
+        package, method_proto.input_type, scope)
+    output_type = self._GetTypeFromScope(
+        package, method_proto.output_type, scope)
+    return descriptor.MethodDescriptor(name=method_proto.name,
+                                       full_name=full_name,
+                                       index=index,
+                                       containing_service=None,
+                                       input_type=input_type,
+                                       output_type=output_type,
+                                       options=_OptionsOrNone(method_proto))
+
   def _ExtractSymbols(self, descriptors):
     """Pulls out all the symbols from descriptor protos.
 
     Args:
       descriptors: The messages to extract descriptors from.
     Yields:
       A two element tuple of the type name and descriptor object.
     """
 
     for desc in descriptors:
@@ skipping 50 matching lines @@
   # TODO(amauryfa): This pool could be constructed from Python code, when we
   # support a flag like 'use_cpp_generated_pool=True'.
   # pylint: disable=protected-access
   _DEFAULT = descriptor._message.default_pool
 else:
   _DEFAULT = DescriptorPool()
 
 
 def Default():
   return _DEFAULT