Skip to content

Commit

Permalink
Allowing get() to accept a single key (in addition to list).
Browse files Browse the repository at this point in the history
  • Loading branch information
dhermes committed Jan 8, 2015
1 parent 6eaa636 commit 1d6cfb4
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 10 deletions.
26 changes: 19 additions & 7 deletions gcloud/datastore/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
Query objects rather than via protobufs.
"""

import collections

from gcloud.datastore import _implicit_environ
from gcloud.datastore import helpers

Expand Down Expand Up @@ -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
Expand All @@ -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 []

Expand Down Expand Up @@ -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):
Expand Down
41 changes: 38 additions & 3 deletions gcloud/datastore/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 1d6cfb4

Please sign in to comment.