diff --git a/gcloud/datastore/dataset.py b/gcloud/datastore/dataset.py index ed5e85ef225a..fb1cc79097c7 100644 --- a/gcloud/datastore/dataset.py +++ b/gcloud/datastore/dataset.py @@ -196,17 +196,5 @@ def allocate_ids(self, incomplete_key, num_ids): self.id(), incomplete_key_pbs) allocated_ids = [allocated_key_pb.path_element[-1].id for allocated_key_pb in allocated_key_pbs] - - # This method is temporary and will move over to Key in - # part 5 of #451. - def create_new_key(new_id): - """Temporary method to complete `incomplete_key`. - - Uses `incomplete_key` from outside scope. - """ - clone = incomplete_key._clone() - clone._path[-1]['id'] = new_id - return clone - - return [create_new_key(allocated_id) + return [incomplete_key.completed_key(allocated_id) for allocated_id in allocated_ids] diff --git a/gcloud/datastore/key.py b/gcloud/datastore/key.py index c731142fb212..6c95439665f3 100644 --- a/gcloud/datastore/key.py +++ b/gcloud/datastore/key.py @@ -127,6 +127,32 @@ def _clone(self): """ return copy.deepcopy(self) + def completed_key(self, id_or_name): + """Creates new key from existing partial key by adding final ID/name. + + :rtype: :class:`gcloud.datastore.key.Key` + :returns: A new `Key` instance with the same data as the current one + and an extra ID or name added. + :raises: `ValueError` if the current key is not partial or if + `id_or_name` is not a string or integer. + """ + if not self.is_partial: + raise ValueError('Only a partial key can be completed.') + + id_or_name_key = None + if isinstance(id_or_name, six.string_types): + id_or_name_key = 'name' + elif isinstance(id_or_name, six.integer_types): + id_or_name_key = 'id' + else: + raise ValueError(id_or_name, + 'ID/name was not a string or integer.') + + new_key = self._clone() + new_key._path[-1][id_or_name_key] = id_or_name + new_key._flat_path += (id_or_name,) + return new_key + def to_protobuf(self): """Return a protobuf corresponding to the key. diff --git a/gcloud/datastore/test_key.py b/gcloud/datastore/test_key.py index 08b79cfa5042..c9309a420d6b 100644 --- a/gcloud/datastore/test_key.py +++ b/gcloud/datastore/test_key.py @@ -62,6 +62,30 @@ def test__clone(self): self.assertEqual(clone.kind, _KIND) self.assertEqual(clone.path, _PATH) + def test_completed_key_on_partial_w_id(self): + key = self._makeOne('KIND') + _ID = 1234 + new_key = key.completed_key(_ID) + self.assertFalse(key is new_key) + self.assertEqual(new_key.id, _ID) + self.assertEqual(new_key.name, None) + + def test_completed_key_on_partial_w_name(self): + key = self._makeOne('KIND') + _NAME = 'NAME' + new_key = key.completed_key(_NAME) + self.assertFalse(key is new_key) + self.assertEqual(new_key.id, None) + self.assertEqual(new_key.name, _NAME) + + def test_completed_key_on_partial_w_invalid(self): + key = self._makeOne('KIND') + self.assertRaises(ValueError, key.completed_key, object()) + + def test_completed_key_on_complete(self): + key = self._makeOne('KIND', 1234) + self.assertRaises(ValueError, key.completed_key, 5678) + def test_to_protobuf_defaults(self): from gcloud.datastore.datastore_v1_pb2 import Key as KeyPB _KIND = 'KIND' diff --git a/regression/datastore.py b/regression/datastore.py index f2ba27d19af8..87b23576c4e6 100644 --- a/regression/datastore.py +++ b/regression/datastore.py @@ -74,9 +74,9 @@ def _get_post(self, name=None, key_id=None, post_content=None): # Update the entity key. key = None if name is not None: - key = datastore.key.Key('Post', name) + key = entity.key().completed_key(name) if key_id is not None: - key = datastore.key.Key('Post', key_id) + key = entity.key().completed_key(key_id) if key is not None: entity.key(key)