diff --git a/gcloud/datastore/api.py b/gcloud/datastore/api.py index 7c0e893e160a..1d54c5d062d4 100644 --- a/gcloud/datastore/api.py +++ b/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): diff --git a/gcloud/datastore/test_api.py b/gcloud/datastore/test_api.py index 1e66b07de181..77ad3122c6ec 100644 --- a/gcloud/datastore/test_api.py +++ b/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