diff --git a/google/cloud/datastore/helpers.py b/google/cloud/datastore/helpers.py index 6eaa3b89..d491360c 100644 --- a/google/cloud/datastore/helpers.py +++ b/google/cloud/datastore/helpers.py @@ -18,7 +18,6 @@ """ import datetime -import itertools from google.protobuf import struct_pb2 from google.type import latlng_pb2 @@ -43,36 +42,29 @@ def _get_meaning(value_pb, is_list=False): :param is_list: Boolean indicating if the ``value_pb`` contains a list value. - :rtype: int + :rtype: int | Tuple[Optional[int], Optional[list[int | None]]] | None :returns: The meaning for the ``value_pb`` if one is set, else - :data:`None`. For a list value, if there are disagreeing - means it just returns a list of meanings. If all the - list meanings agree, it just condenses them. + :data:`None`. For a list value, returns a tuple of + the root meaning of the list, and a list of meanings + of each sub-value. If subvalues are all empty, returns + :data:`None` instead of a list. """ if is_list: + root_meaning = value_pb.meaning or None values = value_pb.array_value.values - # An empty list will have no values, hence no shared meaning - # set among them. - if len(values) == 0: - return None - # We check among all the meanings, some of which may be None, # the rest which may be enum/int values. - all_meanings = [_get_meaning(sub_value_pb) for sub_value_pb in values] - unique_meanings = set(all_meanings) - - if len(unique_meanings) == 1: - # If there is a unique meaning, we preserve it. - return unique_meanings.pop() - else: # We know len(value_pb.array_value.values) > 0. - # If the meaning is not unique, just return all of them. - return all_meanings - - elif value_pb.meaning: # Simple field (int32). - return value_pb.meaning - - return None + sub_meanings = [sub_value_pb.meaning or None for sub_value_pb in values] + if not any(meaning is not None for meaning in sub_meanings): + sub_meanings = None + if root_meaning is None and sub_meanings is None: + # no meanings to save + return None + else: + return root_meaning, sub_meanings + else: + return value_pb.meaning or None def _new_value_pb(entity_pb, name): @@ -156,6 +148,10 @@ def entity_from_protobuf(pb): def _set_pb_meaning_from_entity(entity, name, value, value_pb, is_list=False): """Add meaning information (from an entity) to a protobuf. + value_pb is assumed to have no `meaning` data currently present. + This means if the entity's meaning data is None, this function will do nothing, + rather than removing any existing data. + :type entity: :class:`google.cloud.datastore.entity.Entity` :param entity: The entity to be turned into a protobuf. @@ -181,14 +177,20 @@ def _set_pb_meaning_from_entity(entity, name, value, value_pb, is_list=False): if orig_value is not value: return - # For lists, we set meaning on each sub-element. - if is_list: - if not isinstance(meaning, list): - meaning = itertools.repeat(meaning) - val_iter = zip(value_pb.array_value.values, meaning) - for sub_value_pb, sub_meaning in val_iter: - if sub_meaning is not None: - sub_value_pb.meaning = sub_meaning + if meaning is None: + # no meaning data to set + return + elif is_list: + # for lists, set meaning on the root pb and on each sub-element + root_meaning, sub_meaning_list = meaning + if root_meaning is not None: + value_pb.meaning = root_meaning + if sub_meaning_list: + for sub_value_pb, sub_meaning in zip( + value_pb.array_value.values, sub_meaning_list + ): + if sub_meaning is not None: + sub_value_pb.meaning = sub_meaning else: value_pb.meaning = meaning diff --git a/tests/unit/test_helpers.py b/tests/unit/test_helpers.py index 38702dba..a6f63a80 100644 --- a/tests/unit/test_helpers.py +++ b/tests/unit/test_helpers.py @@ -361,19 +361,21 @@ def test_entity_to_protobuf_w_variable_meanings(): entity = Entity() name = "quux" entity[name] = values = [1, 20, 300] - meaning = 9 - entity._meanings[name] = ([None, meaning, None], values) + root_meaning = 31 + sub_meaning = 9 + entity._meanings[name] = ((root_meaning, [None, sub_meaning, None]), values) entity_pb = entity_to_protobuf(entity) # Construct the expected protobuf. expected_pb = entity_pb2.Entity() value_pb = _new_value_pb(expected_pb, name) + value_pb.meaning = root_meaning value0 = value_pb.array_value.values.add() value0.integer_value = values[0] # The only array entry with a meaning is the middle one. value1 = value_pb.array_value.values.add() value1.integer_value = values[1] - value1.meaning = meaning + value1.meaning = sub_meaning value2 = value_pb.array_value.values.add() value2.integer_value = values[2] @@ -1179,7 +1181,46 @@ def test__get_meaning_w_array_value(): sub_value_pb2.string_value = "bye" result = _get_meaning(value_pb, is_list=True) - assert meaning == result + # should preserve sub-value meanings as list + assert (None, [meaning, meaning]) == result + + +def test__get_meaning_w_array_value_root_meaning(): + from google.cloud.datastore_v1.types import entity as entity_pb2 + from google.cloud.datastore.helpers import _get_meaning + + value_pb = entity_pb2.Value() + meaning = 9 + value_pb.meaning = meaning + sub_value_pb1 = value_pb._pb.array_value.values.add() + sub_value_pb2 = value_pb._pb.array_value.values.add() + + sub_value_pb1.string_value = "hi" + sub_value_pb2.string_value = "bye" + + result = _get_meaning(value_pb, is_list=True) + # should preserve sub-value meanings as list + assert (meaning, None) == result + + +def test__get_meaning_w_array_value_root_and_sub_meanings(): + from google.cloud.datastore_v1.types import entity as entity_pb2 + from google.cloud.datastore.helpers import _get_meaning + + value_pb = entity_pb2.Value() + root_meaning = 9 + sub_meaning = 3 + value_pb.meaning = root_meaning + sub_value_pb1 = value_pb._pb.array_value.values.add() + sub_value_pb2 = value_pb._pb.array_value.values.add() + + sub_value_pb1.meaning = sub_value_pb2.meaning = sub_meaning + sub_value_pb1.string_value = "hi" + sub_value_pb2.string_value = "bye" + + result = _get_meaning(value_pb, is_list=True) + # should preserve sub-value meanings as list + assert (root_meaning, [sub_meaning, sub_meaning]) == result def test__get_meaning_w_array_value_multiple_meanings(): @@ -1198,7 +1239,7 @@ def test__get_meaning_w_array_value_multiple_meanings(): sub_value_pb2.string_value = "bye" result = _get_meaning(value_pb, is_list=True) - assert result == [meaning1, meaning2] + assert result == (None, [meaning1, meaning2]) def test__get_meaning_w_array_value_meaning_partially_unset(): @@ -1215,7 +1256,102 @@ def test__get_meaning_w_array_value_meaning_partially_unset(): sub_value_pb2.string_value = "bye" result = _get_meaning(value_pb, is_list=True) - assert result == [meaning1, None] + assert result == (None, [meaning1, None]) + + +def test__get_meaning_w_array_value_meaning_fully_unset(): + from google.cloud.datastore_v1.types import entity as entity_pb2 + from google.cloud.datastore.helpers import _get_meaning + + value_pb = entity_pb2.Value() + sub_value_pb1 = value_pb._pb.array_value.values.add() + sub_value_pb2 = value_pb._pb.array_value.values.add() + + sub_value_pb1.string_value = "hi" + sub_value_pb2.string_value = "bye" + + result = _get_meaning(value_pb, is_list=True) + assert result is None + + +@pytest.mark.parametrize("orig_root_meaning", [0, 1]) +@pytest.mark.parametrize("orig_sub_meaning", [0, 1]) +def test__set_pb_meaning_w_array_value_fully_unset(orig_root_meaning, orig_sub_meaning): + """ + call _set_pb_meaning_from_entity with meaning=None data. + Should not touch proto's meaning field + """ + from google.cloud.datastore_v1.types import entity as entity_pb2 + from google.cloud.datastore.helpers import _set_pb_meaning_from_entity + from google.cloud.datastore.entity import Entity + + orig_pb = entity_pb2.Entity() + value_pb = orig_pb._pb.properties.get_or_create("value") + value_pb.meaning = orig_root_meaning + sub_value_pb1 = value_pb.array_value.values.add() + sub_value_pb1.meaning = orig_sub_meaning + + entity = Entity(key="key") + entity._meanings = {"value": ((None, None), None)} + _set_pb_meaning_from_entity(entity, "value", None, value_pb, is_list=True) + assert value_pb.meaning == orig_root_meaning + assert value_pb.array_value.values[0].meaning == orig_sub_meaning + + +@pytest.mark.parametrize("orig_meaning", [0, 1]) +def test__set_pb_meaning_w_value_unset(orig_meaning): + """ + call _set_pb_meaning_from_entity with meaning=None data. + Should not touch proto's meaning field + """ + from google.cloud.datastore_v1.types import entity as entity_pb2 + from google.cloud.datastore.helpers import _set_pb_meaning_from_entity + from google.cloud.datastore.entity import Entity + + orig_pb = entity_pb2.Entity() + value_pb = orig_pb._pb.properties.get_or_create("value") + value_pb.meaning = orig_meaning + + entity = Entity(key="key") + entity._meanings = {"value": (None, None)} + _set_pb_meaning_from_entity(entity, "value", None, value_pb, is_list=False) + assert value_pb.meaning == orig_meaning + + +def test__array_w_meaning_end_to_end(): + """ + Test proto->entity->proto with an array with a meaning field + """ + from google.cloud.datastore_v1.types import entity as entity_pb2 + from google.cloud.datastore.helpers import entity_from_protobuf + from google.cloud.datastore.helpers import entity_to_protobuf + + orig_pb = entity_pb2.Entity() + value_pb = orig_pb._pb.properties.get_or_create("value") + value_pb.meaning = 31 + sub_value_pb1 = value_pb.array_value.values.add() + sub_value_pb1.double_value = 1 + sub_value_pb1.meaning = 1 + sub_value_pb2 = value_pb.array_value.values.add() + sub_value_pb2.double_value = 2 + sub_value_pb3 = value_pb.array_value.values.add() + sub_value_pb3.double_value = 3 + sub_value_pb3.meaning = 3 + # convert to entity + entity = entity_from_protobuf(orig_pb._pb) + assert entity._meanings["value"][0] == (31, [1, None, 3]) + assert entity._meanings["value"][1] == [1, 2, 3] + # convert back to pb + output_entity_pb = entity_to_protobuf(entity) + final_pb = output_entity_pb._pb.properties["value"] + assert final_pb.meaning == 31 + assert len(final_pb.array_value.values) == 3 + assert final_pb.array_value.values[0].meaning == 1 + assert final_pb.array_value.values[0].double_value == 1 + assert final_pb.array_value.values[1].meaning == 0 + assert final_pb.array_value.values[1].double_value == 2 + assert final_pb.array_value.values[2].meaning == 3 + assert final_pb.array_value.values[2].double_value == 3 def _make_geopoint(*args, **kwargs):