def _get_methods(
self,
methods: Sequence[descriptor_pb2.MethodDescriptorProto],
address: metadata.Address,
path: Tuple[int, ...],
) -> Mapping[str, wrappers.Method]:
"""Return a dictionary of wrapped methods for the given service.
Args:
methods (Sequence[~.descriptor_pb2.MethodDescriptorProto]): A
sequence of protobuf method objects.
address (~.metadata.Address): An address object denoting the
location of these methods.
path (Tuple[int]): The source location path thus far, as understood
by ``SourceCodeInfo.Location``.
Returns:
Mapping[str, ~.wrappers.Method]: A ordered mapping of
:class:`~.wrappers.Method` objects.
"""
# Iterate over the methods and collect them into a dictionary.
answer: Dict[str, wrappers.Method] = collections.OrderedDict()
for meth_pb, i in zip(methods, range(0, sys.maxsize)):
lro = None
# If the output type is google.longrunning.Operation, we use
# a specialized object in its place.
if meth_pb.output_type.endswith('google.longrunning.Operation'):
op = meth_pb.options.Extensions[operations_pb2.operation_info]
if not op.response_type or not op.metadata_type:
raise TypeError(
f'rpc {meth_pb.name} returns a google.longrunning.'
'Operation, but is missing a response type or '
'metadata type.', )
lro = wrappers.OperationInfo(
response_type=self.api_messages[address.resolve(
op.response_type, )],
metadata_type=self.api_messages[address.resolve(
op.metadata_type, )],
)
# Create the method wrapper object.
answer[meth_pb.name] = wrappers.Method(
input=self.api_messages[meth_pb.input_type.lstrip('.')],
lro=lro,
method_pb=meth_pb,
meta=metadata.Metadata(
address=address.child(meth_pb.name, path + (i, )),
documentation=self.docs.get(path + (i, ), self.EMPTY),
),
output=self.api_messages[meth_pb.output_type.lstrip('.')],
)
# Done; return the answer.
return answer
def _get_methods(
self,
methods: List[descriptor_pb2.MethodDescriptorProto],
address: metadata.Address,
path: Tuple[int],
) -> Mapping[str, wrappers.Method]:
"""Return a dictionary of wrapped methods for the given service.
Args:
methods (Sequence[~.descriptor_pb2.MethodDescriptorProto]): A
sequence of protobuf method objects.
address (~.metadata.Address): An address object denoting the
location of these methods.
path (Tuple[int]): The source location path thus far, as understood
by ``SourceCodeInfo.Location``.
Returns:
Mapping[str, ~.wrappers.Method]: A ordered mapping of
:class:`~.wrappers.Method` objects.
"""
# Iterate over the methods and collect them into a dictionary.
answer = collections.OrderedDict()
for meth_pb, i in zip(methods, range(0, sys.maxsize)):
types = meth_pb.options.Extensions[operations_pb2.operation_types]
# If the output type is google.longrunning.Operation, we use
# a specialized object in its place.
output_type = self.all_messages[meth_pb.output_type.lstrip('.')]
if meth_pb.output_type.endswith('google.longrunning.Operation'):
output_type = self._get_operation_type(
response_type=self.all_messages[address.resolve(
types.response)],
metadata_type=self.all_messages.get(
address.resolve(types.metadata), ),
)
# Create the method wrapper object.
answer[meth_pb.name] = wrappers.Method(
input=self.all_messages[meth_pb.input_type.lstrip('.')],
method_pb=meth_pb,
meta=metadata.Metadata(
address=address,
documentation=self.docs.get(path + (i, ), self.EMPTY),
),
output=output_type,
)
# Done; return the answer.
return answer
def _load_service(
self,
service: descriptor_pb2.ServiceDescriptorProto,
address: metadata.Address,
path: Tuple[int],
) -> wrappers.Service:
"""Load comments for a service and its methods."""
address = address.child(service.name, path)
# Put together a dictionary of the service's methods.
methods = self._get_methods(
service.method,
service_address=address,
path=path + (2, ),
)
# Load the comments for the service itself.
self.proto_services[address.proto] = wrappers.Service(
meta=metadata.Metadata(
address=address,
documentation=self.docs.get(path, self.EMPTY),
),
methods=methods,
service_pb=service,
)
return self.proto_services[address.proto]
def _load_enum(
self,
enum: descriptor_pb2.EnumDescriptorProto,
address: metadata.Address,
path: Tuple[int],
) -> wrappers.EnumType:
"""Load enum descriptions from EnumDescriptorProtos."""
address = address.child(enum.name, path)
# Put together wrapped objects for the enum values.
values = []
for enum_value, i in zip(enum.value, range(0, sys.maxsize)):
values.append(
wrappers.EnumValueType(
enum_value_pb=enum_value,
meta=metadata.Metadata(
address=address,
documentation=self.docs.get(path + (2, i), self.EMPTY),
),
))
# Load the enum itself.
self.proto_enums[address.proto] = wrappers.EnumType(
enum_pb=enum,
meta=metadata.Metadata(
address=address,
documentation=self.docs.get(path, self.EMPTY),
),
values=values,
)
return self.proto_enums[address.proto]
def _load_message(
self,
message_pb: descriptor_pb2.DescriptorProto,
address: metadata.Address,
path: Tuple[int],
) -> wrappers.MessageType:
"""Load message descriptions from DescriptorProtos."""
address = address.child(message_pb.name, path)
# Load all nested items.
#
# Note: This occurs before piecing together this message's fields
# because if nested types are present, they are generally the
# type of one of this message's fields, and they need to be in
# the registry for the field's message or enum attributes to be
# set correctly.
nested_enums = self._load_children(
message_pb.enum_type,
address=address,
loader=self._load_enum,
path=path + (4, ),
)
nested_messages = self._load_children(
message_pb.nested_type,
address=address,
loader=self._load_message,
path=path + (3, ),
)
# self._load_children(message.oneof_decl, loader=self._load_field,
# address=nested_addr, info=info.get(8, {}))
# Create a dictionary of all the fields for this message.
fields = self._get_fields(
message_pb.field,
address=address,
path=path + (2, ),
)
fields.update(
self._get_fields(
message_pb.extension,
address=address,
path=path + (6, ),
))
# Create a message correspoding to this descriptor.
self.proto_messages[address.proto] = wrappers.MessageType(
fields=fields,
message_pb=message_pb,
nested_enums=nested_enums,
nested_messages=nested_messages,
meta=metadata.Metadata(
address=address,
documentation=self.docs.get(path, self.EMPTY),
),
)
return self.proto_messages[address.proto]
def _get_fields(
self,
field_pbs: Sequence[descriptor_pb2.FieldDescriptorProto],
address: metadata.Address,
path: Tuple[int, ...],
oneofs: Optional[Dict[str, wrappers.Oneof]] = None
) -> Dict[str, wrappers.Field]:
"""Return a dictionary of wrapped fields for the given message.
Args:
field_pbs (Sequence[~.descriptor_pb2.FieldDescriptorProto]): A
sequence of protobuf field objects.
address (~.metadata.Address): An address object denoting the
location of these fields.
path (Tuple[int]): The source location path thus far, as
understood by ``SourceCodeInfo.Location``.
Returns:
Mapping[str, ~.wrappers.Field]: A ordered mapping of
:class:`~.wrappers.Field` objects.
"""
# Iterate over the fields and collect them into a dictionary.
#
# The saving of the enum and message types rely on protocol buffers'
# naming rules to trust that they will never collide.
#
# Note: If this field is a recursive reference to its own message,
# then the message will not be in `api_messages` yet (because the
# message wrapper is not yet created, because it needs this object
# first) and this will be None. This case is addressed in the
# `_load_message` method.
answer: Dict[str, wrappers.Field] = collections.OrderedDict()
for i, field_pb in enumerate(field_pbs):
is_oneof = oneofs and field_pb.HasField('oneof_index')
oneof_name = nth((oneofs or {}).keys(),
field_pb.oneof_index) if is_oneof else None
field = wrappers.Field(
field_pb=field_pb,
enum=self.api_enums.get(field_pb.type_name.lstrip('.')),
message=self.api_messages.get(field_pb.type_name.lstrip('.')),
meta=metadata.Metadata(
address=address.child(field_pb.name, path + (i, )),
documentation=self.docs.get(path + (i, ), self.EMPTY),
),
oneof=oneof_name,
)
answer[field.name] = field
# Done; return the answer.
return answer
def _get_methods(self,
methods: Sequence[descriptor_pb2.MethodDescriptorProto],
service_address: metadata.Address, path: Tuple[int, ...],
) -> Mapping[str, wrappers.Method]:
"""Return a dictionary of wrapped methods for the given service.
Args:
methods (Sequence[~.descriptor_pb2.MethodDescriptorProto]): A
sequence of protobuf method objects.
service_address (~.metadata.Address): An address object for the
service, denoting the location of these methods.
path (Tuple[int]): The source location path thus far, as understood
by ``SourceCodeInfo.Location``.
Returns:
Mapping[str, ~.wrappers.Method]: A ordered mapping of
:class:`~.wrappers.Method` objects.
"""
# Iterate over the methods and collect them into a dictionary.
answer: Dict[str, wrappers.Method] = collections.OrderedDict()
for i, meth_pb in enumerate(methods):
retry, timeout = self._get_retry_and_timeout(
service_address,
meth_pb
)
# Create the method wrapper object.
answer[meth_pb.name] = wrappers.Method(
input=self.api_messages[meth_pb.input_type.lstrip('.')],
lro=self._maybe_get_lro(service_address, meth_pb),
extended_lro=self._maybe_get_extended_lro(
service_address,
meth_pb,
),
method_pb=meth_pb,
meta=metadata.Metadata(
address=service_address.child(meth_pb.name, path + (i,)),
documentation=self.docs.get(path + (i,), self.EMPTY),
),
output=self.api_messages[meth_pb.output_type.lstrip('.')],
retry=retry,
timeout=timeout,
)
# Done; return the answer.
return answer
def _maybe_get_lro(
self,
service_address: metadata.Address,
meth_pb: descriptor_pb2.MethodDescriptorProto
) -> Optional[wrappers.OperationInfo]:
"""Determines whether a method is a Long Running Operation (aka LRO)
and, if it is, return an OperationInfo that includes the response
and metadata types.
Args:
service_address (~.metadata.Address): An address object for the
service, denoting the location of these methods.
meth_pb (~.descriptor_pb2.MethodDescriptorProto): A
protobuf method objects.
Returns:
Optional[~.wrappers.OperationInfo]: The info for the long-running
operation, if the passed method is an LRO.
"""
lro = None
# If the output type is google.longrunning.Operation, we use
# a specialized object in its place.
if meth_pb.output_type.endswith('google.longrunning.Operation'):
if not meth_pb.options.HasExtension(operations_pb2.operation_info):
# This is not a long running operation even though it returns
# an Operation.
return None
op = meth_pb.options.Extensions[operations_pb2.operation_info]
if not op.response_type or not op.metadata_type:
raise TypeError(
f'rpc {meth_pb.name} returns a google.longrunning.'
'Operation, but is missing a response type or '
'metadata type.',
)
response_key = service_address.resolve(op.response_type)
metadata_key = service_address.resolve(op.metadata_type)
lro = wrappers.OperationInfo(
response_type=self.api_messages[response_key],
metadata_type=self.api_messages[metadata_key],
)
return lro
def _get_methods(
self,
methods: Sequence[descriptor_pb2.MethodDescriptorProto],
service_address: metadata.Address,
path: Tuple[int, ...],
) -> Mapping[str, wrappers.Method]:
"""Return a dictionary of wrapped methods for the given service.
Args:
methods (Sequence[~.descriptor_pb2.MethodDescriptorProto]): A
sequence of protobuf method objects.
service_address (~.metadata.Address): An address object for the
service, denoting the location of these methods.
path (Tuple[int]): The source location path thus far, as understood
by ``SourceCodeInfo.Location``.
Returns:
Mapping[str, ~.wrappers.Method]: A ordered mapping of
:class:`~.wrappers.Method` objects.
"""
# Iterate over the methods and collect them into a dictionary.
answer: Dict[str, wrappers.Method] = collections.OrderedDict()
for meth_pb, i in zip(methods, range(0, sys.maxsize)):
lro = None
# If the output type is google.longrunning.Operation, we use
# a specialized object in its place.
if meth_pb.output_type.endswith('google.longrunning.Operation'):
op = meth_pb.options.Extensions[operations_pb2.operation_info]
if not op.response_type or not op.metadata_type:
raise TypeError(
f'rpc {meth_pb.name} returns a google.longrunning.'
'Operation, but is missing a response type or '
'metadata type.', )
lro = wrappers.OperationInfo(
response_type=self.api_messages[service_address.resolve(
op.response_type, )],
metadata_type=self.api_messages[service_address.resolve(
op.metadata_type, )],
)
# If we got a gRPC service config, get the appropriate retry
# and timeout information from it.
retry = None
timeout = None
# This object should be a dictionary that conforms to the
# gRPC service config proto:
# Repo: https://github.com/grpc/grpc-proto/
# Filename: grpc/service_config/service_config.proto
#
# We only care about a small piece, so we are just leaving
# it as a dictionary and parsing accordingly.
if self.opts.retry:
# The gRPC service config uses a repeated `name` field
# with a particular format, which we match against.
# This defines the expected selector for *this* method.
selector = {
'service':
'{package}.{service_name}'.format(
package='.'.join(service_address.package),
service_name=service_address.name,
),
'method':
meth_pb.name,
}
# Find the method config that applies to us, if any.
mc = next((i for i in self.opts.retry.get('methodConfig', [])
if selector in i.get('name')), None)
if mc:
# Set the timeout according to this method config.
if mc.get('timeout'):
timeout = self._to_float(mc['timeout'])
# Set the retry according to this method config.
if 'retryPolicy' in mc:
r = mc['retryPolicy']
retry = wrappers.RetryInfo(
max_attempts=r.get('maxAttempts', 0),
initial_backoff=self._to_float(
r.get('initialBackoff', '0s'), ),
max_backoff=self._to_float(
r.get('maxBackoff', '0s'), ),
backoff_multiplier=r.get('backoffMultiplier', 0.0),
retryable_exceptions=frozenset(
exceptions.exception_class_for_grpc_status(
getattr(grpc.StatusCode, code), )
for code in r.get('retryableStatusCodes', [])),
)
# Create the method wrapper object.
answer[meth_pb.name] = wrappers.Method(
input=self.api_messages[meth_pb.input_type.lstrip('.')],
lro=lro,
method_pb=meth_pb,
meta=metadata.Metadata(
address=service_address.child(meth_pb.name, path + (i, )),
documentation=self.docs.get(path + (i, ), self.EMPTY),
),
output=self.api_messages[meth_pb.output_type.lstrip('.')],
retry=retry,
timeout=timeout,
)
# Done; return the answer.
return answer
def _maybe_get_extended_lro(
self,
service_address: metadata.Address,
meth_pb: descriptor_pb2.MethodDescriptorProto,
) -> Optional[wrappers.ExtendedOperationInfo]:
op_service_name = meth_pb.options.Extensions[ex_ops_pb2.operation_service]
if not op_service_name:
return None
# Manual lookups because services and methods haven't been loaded.
# Note: this assumes that the operation service lives in the same proto file.
# This is a reasonable assumption as of March '22, but it may (unlikely)
# change in the future.
op_service_pb = next(
(s for s in self.file_descriptor.service if s.name == op_service_name),
None,
)
if not op_service_pb:
raise ValueError(
f"Could not find custom operation service: {op_service_name}"
)
operation_polling_method_pb = next(
(
m
for m in op_service_pb.method
if m.options.Extensions[ex_ops_pb2.operation_polling_method]
),
None,
)
if not operation_polling_method_pb:
raise ValueError(
f"Could not find operation polling method for custom operation service: {op_service_name}"
)
operation_request_key = service_address.resolve(
operation_polling_method_pb.input_type.lstrip(".")
)
operation_request_message = self.api_messages[operation_request_key]
operation_type = service_address.resolve(
operation_polling_method_pb.output_type.lstrip(".")
)
method_output_type = service_address.resolve(
meth_pb.output_type.lstrip(".")
)
if operation_type != method_output_type:
raise ValueError(
f"Inconsistent return types between extended lro method '{meth_pb.name}'"
f" and extended lro polling method '{operation_polling_method_pb.name}':"
f" '{method_output_type}' and '{operation_type}'"
)
operation_message = self.api_messages[operation_type]
if not operation_message.is_extended_operation:
raise ValueError(
f"Message is not an extended operation: {operation_type}"
)
return wrappers.ExtendedOperationInfo(
request_type=operation_request_message,
operation_type=operation_message,
)
