Implement Key.complete_key for auto_id on partial keys.

dhermes · dhermes · commit 33d546d61eb1 · 2014-12-31T13:12:55.000-08:00
Addresses fifth part of #451.
diff --git 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.complete_key(allocated_id)
                 for allocated_id in allocated_ids]
diff --git a/gcloud/datastore/key.py b/gcloud/datastore/key.py
@@ -127,6 +127,32 @@ def _clone(self):
         """
         return copy.deepcopy(self)
 
+    def complete_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
@@ -62,6 +62,30 @@ def test__clone(self):
         self.assertEqual(clone.kind, _KIND)
         self.assertEqual(clone.path, _PATH)
 
+    def test_complete_key_on_partial_w_id(self):
+        key = self._makeOne('KIND')
+        _ID = 1234
+        new_key = key.complete_key(_ID)
+        self.assertFalse(key is new_key)
+        self.assertEqual(new_key.id, _ID)
+        self.assertEqual(new_key.name, None)
+
+    def test_complete_key_on_partial_w_name(self):
+        key = self._makeOne('KIND')
+        _NAME = 'NAME'
+        new_key = key.complete_key(_NAME)
+        self.assertFalse(key is new_key)
+        self.assertEqual(new_key.id, None)
+        self.assertEqual(new_key.name, _NAME)
+
+    def test_complete_key_on_partial_w_invalid(self):
+        key = self._makeOne('KIND')
+        self.assertRaises(ValueError, key.complete_key, object())
+
+    def test_complete_key_on_complete(self):
+        key = self._makeOne('KIND', 1234)
+        self.assertRaises(ValueError, key.complete_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
@@ -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().complete_key(name)
         if key_id is not None:
-            key = datastore.key.Key('Post', key_id)
+            key = entity.key().complete_key(key_id)
         if key is not None:
             entity.key(key)
 
