Adding a non-protobuf allocate_ids method on Dataset.

gcloud/datastore/dataset.py

@@ -182,3 +182,29 @@ def get_entities(self, keys, missing=None, deferred=None):
             entities.append(helpers.entity_from_protobuf(
                 entity_pb, dataset=self))
         return entities
+
+    def allocate_ids(self, incomplete_key, num_ids):
+        """Allocates a list of IDs from a partial key.
+
+        :type incomplete_key: A :class:`gcloud.datastore.key.Key`
+        :param incomplete_key: Partial key to use as base for allocated IDs.
+
+        :type num_ids: A :class:`int`.
+        :param num_ids: The number of IDs to allocate.
+
+        :rtype: list of :class:`gcloud.datastore.key.Key`
+        :return: The (complete) keys allocated with `incomplete_key` as root.
+        :raises: `ValueError` if `incomplete_key` is not a partial key.
+        """
+        if not incomplete_key.is_partial():
+            raise ValueError(('Key is not partial.', incomplete_key))
+
+        incomplete_key_pb = incomplete_key.to_protobuf()
+        incomplete_key_pbs = [incomplete_key_pb] * num_ids
+
+        allocated_key_pbs = self.connection().allocate_ids(
+            self.id(), incomplete_key_pbs)
+        allocated_ids = [allocated_key_pb.path_element[-1].id
+                         for allocated_key_pb in allocated_key_pbs]
+        return [incomplete_key.id(allocated_id)
+                for allocated_id in allocated_ids]

gcloud/datastore/test_dataset.py

@@ -213,6 +213,40 @@ def test_get_entities_hit(self):
         self.assertEqual(list(result), ['foo'])
         self.assertEqual(result['foo'], 'Foo')
 
+    def test_allocate_ids(self):
+        from gcloud.datastore.test_entity import _Key
+
+        INCOMPLETE_KEY = _Key()
+        PROTO_ID = object()
+        INCOMPLETE_KEY._key = _KeyProto(PROTO_ID)
+        INCOMPLETE_KEY._partial = True
+
+        CONNECTION = _Connection()
+        NUM_IDS = 2
+        DATASET_ID = 'foo'
+        DATASET = self._makeOne(DATASET_ID, connection=CONNECTION)
+        result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
+
+        # Check the IDs returned match _PathElementProto.
+        self.assertEqual([key._id for key in result], range(NUM_IDS))
+
+        # Check connection is called correctly.
+        self.assertEqual(CONNECTION._called_dataset_id, DATASET_ID)
+        self.assertEqual(len(CONNECTION._called_key_pbs), NUM_IDS)
+
+        # Check the IDs passed to Connection.allocate_ids.
+        key_paths = [key_pb.path_element[-1].id
+                     for key_pb in CONNECTION._called_key_pbs]
+        self.assertEqual(key_paths, [PROTO_ID] * NUM_IDS)
+
+    def test_allocate_ids_with_complete(self):
+        from gcloud.datastore.test_entity import _Key
+
+        COMPLETE_KEY = _Key()
+        DATASET = self._makeOne(None)
+        self.assertRaises(ValueError, DATASET.allocate_ids,
+                          COMPLETE_KEY, 2)
+
 
 class _Connection(object):
     _called_with = None
@@ -230,3 +264,21 @@ def lookup(self, **kw):
         if deferred is not None:
             deferred.extend(self._deferred)
         return self._result
+
+    def allocate_ids(self, dataset_id, key_pbs):
+        self._called_dataset_id = dataset_id
+        self._called_key_pbs = key_pbs
+        num_pbs = len(key_pbs)
+        return [_KeyProto(i) for i in range(num_pbs)]
+
+
+class _PathElementProto(object):
+
+    def __init__(self, _id):
+        self.id = _id
+
+
+class _KeyProto(object):
+
+    def __init__(self, id_):
+        self.path_element = [_PathElementProto(id_)]

gcloud/datastore/test_entity.py

@@ -239,6 +239,13 @@ class _Key(object):
     _key = 'KEY'
     _partial = False
     _path = None
+    _id = None
+
+    def id(self, id_to_set):
+        self._called_id = id_to_set
+        clone = _Key()
+        clone._id = id_to_set
+        return clone
 
     def to_protobuf(self):
         return self._key

regression/datastore.py

@@ -42,18 +42,16 @@ class TestDatastoreAllocateIDs(TestDatastore):
 
     def test_allocate_ids(self):
         incomplete_key = datastore.key.Key(path=[{'kind': 'Kind'}])
-        incomplete_key_pb = incomplete_key.to_protobuf()
-        incomplete_key_pbs = [incomplete_key_pb] * 10
-
-        connection = self.dataset.connection()
-        allocated_key_pbs = connection.allocate_ids(self.dataset.id(),
-                                                    incomplete_key_pbs)
-        allocated_keys = [datastore.helpers.key_from_protobuf(key_pb)
-                          for key_pb in allocated_key_pbs]
+        allocated_keys = self.dataset.allocate_ids(incomplete_key, 10)
         self.assertEqual(len(allocated_keys), 10)
+
+        unique_ids = set()
         for key in allocated_keys:
+            unique_ids.add(key.id())
             self.assertFalse(key.is_partial())
 
+        self.assertEqual(len(unique_ids), 10)
+
 
 class TestDatastoreSave(TestDatastore):