Allowing get() to accept a single key (in addition to list).

Author: dhermes

gcloud/datastore/api.py

@@ -18,6 +18,8 @@
 Query objects rather than via protobufs.
 """
 
+import collections
+
 from gcloud.datastore import _implicit_environ
 from gcloud.datastore import helpers
 
@@ -58,11 +60,12 @@ def _require_connection(connection=None):
     return connection
 
 
-def get(keys, missing=None, deferred=None, connection=None):
+def get(key_or_keys, missing=None, deferred=None, connection=None):
     """Retrieves entities, along with their attributes.
 
-    :type keys: list of :class:`gcloud.datastore.key.Key`
-    :param keys: The name of the item to retrieve.
+    :type key_or_keys: list of :class:`gcloud.datastore.key.Key` or single
+                       :class:`gcloud.datastore.key.Key`
+    :param key_or_keys: The key or keys to be retrieved from the datastore.
 
     :type missing: an empty list or None.
     :param missing: If a list is passed, the key-only entities returned
@@ -77,10 +80,15 @@ def get(keys, missing=None, deferred=None, connection=None):
     :type connection: :class:`gcloud.datastore.connection.Connection`
     :param connection: Optional. The connection used to connect to datastore.
 
-    :rtype: list of :class:`gcloud.datastore.entity.Entity`
-    :returns: The requested entities.
-    :raises: :class:`ValueError` if the key dataset IDs don't agree.
+    :rtype: list of :class:`gcloud.datastore.entity.Entity`, single
+            :class:`gcloud.datastore.entity.Entity`, or ``NoneType``
+    :returns: The requested entities, or single entity.
     """
+    if isinstance(key_or_keys, collections.Iterable):
+        keys = key_or_keys
+    else:
+        keys = [key_or_keys]
+
     if not keys:
         return []
 
@@ -114,7 +122,11 @@ def get(keys, missing=None, deferred=None, connection=None):
     for entity_pb in entity_pbs:
         entities.append(helpers.entity_from_protobuf(entity_pb))
 
-    return entities
+    if keys is key_or_keys:
+        return entities
+    else:
+        if entities:
+            return entities[0]
 
 
 def allocate_ids(incomplete_key, num_ids, connection=None):

gcloud/datastore/test_api.py

@@ -101,15 +101,23 @@ def test_get_no_keys(self):
         results = self._callFUT([])
         self.assertEqual(results, [])
 
-    def test_get_miss(self):
+    def _miss_helper(self, expected_results, use_list=True):
         from gcloud.datastore.key import Key
         from gcloud.datastore.test_connection import _Connection
 
         DATASET_ID = 'DATASET'
         connection = _Connection()
         key = Key('Kind', 1234, dataset_id=DATASET_ID)
-        results = self._callFUT([key], connection=connection)
-        self.assertEqual(results, [])
+        if use_list:
+            key = [key]
+        results = self._callFUT(key, connection=connection)
+        self.assertEqual(results, expected_results)
+
+    def test_get_miss(self):
+        self._miss_helper([], use_list=True)
+
+    def test_get_miss_single_key(self):
+        self._miss_helper(None, use_list=False)
 
     def test_get_miss_w_missing(self):
         from gcloud.datastore import datastore_v1_pb2 as datastore_pb
@@ -240,6 +248,33 @@ def test_get_hit_multiple_keys_different_dataset(self):
         with self.assertRaises(ValueError):
             self._callFUT([key1, key2], connection=object())
 
+    def test_get_hit_single_key(self):
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_connection import _Connection
+
+        DATASET_ID = 'DATASET'
+        KIND = 'Kind'
+        ID = 1234
+        PATH = [{'kind': KIND, 'id': ID}]
+
+        # Make a found entity pb to be returned from mock backend.
+        entity_pb = self._make_entity_pb(DATASET_ID, KIND, ID,
+                                         'foo', 'Foo')
+
+        # Make a connection to return the entity pb.
+        connection = _Connection(entity_pb)
+
+        key = Key(KIND, ID, dataset_id=DATASET_ID)
+        result = self._callFUT(key, connection=connection)
+        new_key = result.key
+
+        # Check the returned value is as expected.
+        self.assertFalse(new_key is key)
+        self.assertEqual(new_key.dataset_id, DATASET_ID)
+        self.assertEqual(new_key.path, PATH)
+        self.assertEqual(list(result), ['foo'])
+        self.assertEqual(result['foo'], 'Foo')
+
     def test_get_implicit(self):
         from gcloud.datastore import _implicit_environ
         from gcloud.datastore.key import Key