Example #1
0
def _get_any_parsed_field(value_field, type_url_field, field_name):
    """Helper function for _get_any_parsed_fields."""
    full_name = get_full_name_from_any_step(field_name)
    indices_with_type = _any_indices_with_type(type_url_field, full_name)
    [index_to_solution_index, index_to_values
     ] = struct2tensor_ops.equi_join_indices(indices_with_type,
                                             value_field.index)
    solution_index = tf.gather(indices_with_type, index_to_solution_index)
    solution_value = tf.gather(value_field.value, index_to_values)
    # TODO(martinz): make _ParsedField public.
    return struct2tensor_ops._ParsedField(  # pylint: disable=protected-access
        field_name=field_name,
        field_descriptor=None,
        index=solution_index,
        value=solution_value)
Example #2
0
def _get_map_parsed_fields(
    desc: descriptor.Descriptor,
    regular_fields: Mapping[StrStep, struct2tensor_ops._ParsedField],
    field_names: Set[StrStep],
    backing_str_tensor: Optional[tf.Tensor] = None
) -> Mapping[StrStep, struct2tensor_ops._ParsedField]:
    """Gets the map proto ParsedFields.

  field_names includes all the fields: map fields, any fields, and
  regular fields.

  Args:
    desc: the descriptor of the parent proto.
    regular_fields: the fields that are parsed directly from the proto.
    field_names: all fields needed: map fields, any fields, and regular fields.
    backing_str_tensor: a string tensor representing the root serialized proto.
      This is passed to keep string_views of the tensor valid for all children
      of the root expression

  Returns:
    A map from field names to ParsedFields, only for the field names of the form
    foo[bar].
  """
    maps_to_parse = collections.defaultdict(dict)
    for x in field_names:
        if path.is_map_indexing_step(x):
            map_field_name, key = path.parse_map_indexing_step(x)
            maps_to_parse[map_field_name][key] = x
    result_as_list = []
    for map_field_name, v in maps_to_parse.items():
        parsed_map_field = regular_fields[map_field_name]
        keys_needed = list(v.keys())
        map_field_value = parsed_map_field.value
        map_field_index = parsed_map_field.index
        map_field_desc = desc.fields_by_name[map_field_name].message_type
        values_and_parent_indices = struct2tensor_ops.parse_proto_map(
            map_field_value, map_field_index, map_field_desc, keys_needed,
            backing_str_tensor)
        for map_key, [value, parent_index] in zip(keys_needed,
                                                  values_and_parent_indices):
            result_as_list.append(
                struct2tensor_ops._ParsedField(field_name=v[map_key],
                                               field_descriptor=None,
                                               index=parent_index,
                                               value=value))
    return {x.field_name: x for x in result_as_list}
Example #3
0
def _get_map_parsed_fields(desc, regular_fields, field_names):
    """Gets the map proto ParsedFields.

  field_names includes all the fields: map fields, any fields, and
  regular fields.

  Args:
    desc: the descriptor of the parent proto.
    regular_fields: the fields that are parsed directly from the proto.
    field_names: all fields needed: map fields, any fields, and regular fields.

  Returns:
    A map from field names to ParsedFields, only for the field names of the form
    foo[bar].
  """
    maps_to_parse = collections.defaultdict(dict)
    for x in field_names:
        if path.is_map_indexing_step(x):
            map_field_name, key = path.parse_map_indexing_step(x)
            maps_to_parse[map_field_name][key] = x
    result_as_list = []
    for map_field_name, v in maps_to_parse.items():
        parsed_map_field = regular_fields[map_field_name]
        keys_needed = list(v.keys())
        map_field_value = parsed_map_field.value
        map_field_index = parsed_map_field.index
        map_field_desc = desc.fields_by_name[map_field_name].message_type
        values_and_parent_indices = struct2tensor_ops.parse_proto_map(
            map_field_value, map_field_index, map_field_desc, keys_needed)
        for map_key, [value, parent_index] in zip(keys_needed,
                                                  values_and_parent_indices):
            result_as_list.append(
                struct2tensor_ops._ParsedField(field_name=v[map_key],
                                               field_descriptor=None,
                                               index=parent_index,
                                               value=value))
    return {x.field_name: x for x in result_as_list}