Fix #81 getting a datastore.entity by key or path

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):
@@ -107,16 +108,27 @@ def transaction(self, *args, **kwargs):
         kwargs['dataset'] = self
         return Transaction(*args, **kwargs)
 
-    def get_entity(self, key):
+    def get_entity(self, key_or_path):
         """Retrieves entity from the dataset, along with its attributes.
 
-        :type key: :class:`gcloud.datastore.key.Key`
-        :param item_name: The name of the item to retrieve.
+        :type key_or_path: :class:`gcloud.datastore.key.Key` or path
+        :param key_or_path: The name of the item to retrieve or 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).
 
         :rtype: :class:`gcloud.datastore.entity.Entity` or ``None``
         :return: The requested entity, or ``None`` if there was no match found.
         """
-        entities = self.get_entities([key])
+
+        if isinstance(key_or_path, Key):
+            entities = self.get_entities([key_or_path])
+        else:
+            key = Key.from_path(*key_or_path)
+            entities = self.get_entities([key])
+
         if entities:
             return entities[0]
 

gcloud/datastore/test_dataset.py

@@ -128,6 +128,41 @@ def test_get_entity_hit(self):
         self.assertEqual(list(result), ['foo'])
         self.assertEqual(result['foo'], 'Foo')
 
+    def test_get_entity_path(self):
+        from gcloud.datastore.connection import datastore_pb
+        DATASET_ID = 'DATASET'
+        KIND = 'Kind'
+        ID = 1234
+        PATH = [{'kind': KIND, 'id': ID}]
+        entity_pb = datastore_pb.Entity()
+        entity_pb.key.partition_id.dataset_id = DATASET_ID
+        path_element = entity_pb.key.path_element.add()
+        path_element.kind = KIND
+        path_element.id = ID
+        prop = entity_pb.property.add()
+        prop.name = 'foo'
+        prop.value.string_value = 'Foo'
+        connection = _Connection(entity_pb)
+        dataset = self._makeOne(DATASET_ID, connection)
+        result = dataset.get_entity([KIND, ID])
+        key = result.key()
+        self.assertEqual(key._dataset_id, DATASET_ID)
+        self.assertEqual(key.path(), PATH)
+        self.assertEqual(list(result), ['foo'])
+        self.assertEqual(result['foo'], 'Foo')
+
+    def test_get_entity_odd_nonetype(self):
+        from gcloud.datastore.connection import datastore_pb
+        DATASET_ID = 'DATASET'
+        KIND = 'Kind'
+        entity_pb = datastore_pb.Entity()
+        connection = _Connection(entity_pb)
+        dataset = self._makeOne(DATASET_ID, connection)
+        with self.assertRaises(ValueError):
+            dataset.get_entity([KIND])
+        with self.assertRaises(TypeError):
+            dataset.get_entity(None)
+
 
 class _Connection(object):
     _called_with = None