Fix #82 creating a new datastore entity with a name is a pain

gcloud/datastore/dataset.py

@@ -4,6 +4,7 @@
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.query import Query
 from gcloud.datastore.transaction import Transaction
+from gcloud.datastore.key import Key
 
 
 class Dataset(object):
@@ -78,21 +79,45 @@ def query(self, *args, **kwargs):
         kwargs['dataset'] = self
         return Query(*args, **kwargs)
 
-    def entity(self, kind, exclude_from_indexes=()):
+    def entity(self, kind_key_path, exclude_from_indexes=()):
         """Create an entity bound to this dataset.
 
-        :type kind: string
+        :type kind_key_path: string, :class:`glcouddatastore.key.Key`, or path
         :param kind: the "kind" of the new entity (see
                  https://cloud.google.com/datastore/docs/concepts/entities#Datastore_Kinds_and_identifiers)
+                     or put the "key" you want to set on the entity or the
+                     "sequence" of even length, where the first of each pair
+                     is a string representing the 'kind' of the path element,
+                     and the second of the pair is either a string (for the
+                     path element's name) or an integer (for its id).
 
         :param exclude_from_indexes: names of fields whose values are not to
                                      be indexed.
 
         :rtype: :class:`gcloud.datastore.entity.Entity`
         :returns: a new Entity instance, bound to this dataset.
         """
-        return Entity(dataset=self, kind=kind,
-                      exclude_from_indexes=exclude_from_indexes)
+
+        if isinstance(kind_key_path, basestring):
+            return Entity(
+                dataset=self, kind=kind_key_path,
+                exclude_from_indexes=exclude_from_indexes)
+
+        elif isinstance(kind_key_path, Key):
+            kind = kind_key_path.kind()
+            return Entity(
+                dataset=self, kind=kind,
+                exclude_from_indexes=exclude_from_indexes).key(kind_key_path)
+
+        elif isinstance(kind_key_path, list):
+            key = Key.from_path(*kind_key_path)
+            kind = key.kind()
+            return Entity(
+                dataset=self, kind=kind,
+                exclude_from_indexes=exclude_from_indexes).key(key)
+
+        else:
+            raise ValueError('Must pass a kind, key or path.')
 
     def transaction(self, *args, **kwargs):
         """Create a transaction bound to this dataset.

gcloud/datastore/test_dataset.py

@@ -54,6 +54,40 @@ def test_entity_factory_explicit(self):
         self.assertEqual(entity.kind(), KIND)
         self.assertEqual(sorted(entity.exclude_from_indexes()), ['bar', 'foo'])
 
+    def test_entity_factory_key(self):
+        from gcloud.datastore.entity import Entity
+        from gcloud.datastore.key import Key
+        DATASET_ID = 'DATASET'
+        KIND = 'KIND'
+        ID = 1234
+        PATH = [{'kind': KIND, 'id': ID}]
+        dataset = self._makeOne(DATASET_ID)
+        key = Key(path=PATH)
+        entity = dataset.entity(key)
+        self.assertIsInstance(entity, Entity)
+        self.assertEqual(entity.kind(), KIND)
+        self.assertEqual(sorted(entity.exclude_from_indexes()), [])
+        self.assertEqual(entity.key().path(), PATH)
+
+    def test_entity_factory_path(self):
+        from gcloud.datastore.entity import Entity
+        DATASET_ID = 'DATASET'
+        KIND = 'KIND'
+        ID = 1234
+        PATH = [{'kind': KIND, 'id': ID}]
+        dataset = self._makeOne(DATASET_ID)
+        entity = dataset.entity([KIND, ID])
+        self.assertIsInstance(entity, Entity)
+        self.assertEqual(entity.kind(), KIND)
+        self.assertEqual(sorted(entity.exclude_from_indexes()), [])
+        self.assertEqual(entity.key().path(), PATH)
+
+    def test_entity_factory_odd_Nontype(self):
+        DATASET_ID = 'DATASET'
+        dataset = self._makeOne(DATASET_ID)
+        self.assertRaises(ValueError, dataset.entity, ['KIND'])
+        self.assertRaises(ValueError, dataset.entity, None)
+
     def test_transaction_factory(self):
         from gcloud.datastore.transaction import Transaction
         DATASET_ID = 'DATASET'