Merge pull request #582 from dhermes/connection-lookup-only-list

Requiring Connection.lookup to take a list in datastore.
googleapis · Feb 3, 2015 · f854cd1 · f854cd1
2 parents d858ff0 + 57e11b1
commit f854cd1
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 23 deletions.
diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py
@@ -136,20 +136,19 @@ def lookup(self, dataset_id, key_pbs,
         >>> from gcloud import datastore
         >>> datastore.set_defaults()
         >>> key = datastore.Key('MyKind', 1234, dataset_id='dataset-id')
-        >>> datastore.get(key)
-        <Entity object>
+        >>> datastore.get([key])
+        [<Entity object>]
 
         Using the ``connection`` class directly:
 
-        >>> connection.lookup('dataset-id', key.to_protobuf())
-        <Entity protobuf>
+        >>> connection.lookup('dataset-id', [key.to_protobuf()])
+        [<Entity protobuf>]
 
         :type dataset_id: string
         :param dataset_id: The ID of the dataset to look up the keys.
 
         :type key_pbs: list of :class:`gcloud.datastore._datastore_v1_pb2.Key`
-                       (or a single Key)
-        :param key_pbs: The key (or keys) to retrieve from the datastore.
+        :param key_pbs: The keys to retrieve from the datastore.
 
         :type missing: an empty list or None.
         :param missing: If a list is passed, the key-only entity protobufs
@@ -188,12 +187,6 @@ def lookup(self, dataset_id, key_pbs,
 
         lookup_request = datastore_pb.LookupRequest()
         _set_read_options(lookup_request, eventual, transaction_id)
-
-        single_key = isinstance(key_pbs, datastore_pb.Key)
-
-        if single_key:
-            key_pbs = [key_pbs]
-
         helpers._add_keys_to_request(lookup_request.key, key_pbs)
 
         results, missing_found, deferred_found = self._lookup(
@@ -205,12 +198,6 @@ def lookup(self, dataset_id, key_pbs,
         if deferred is not None:
             deferred.extend(deferred_found)
 
-        if single_key:
-            if results:
-                return results[0]
-            else:
-                return None
-
         return results
 
     def run_query(self, dataset_id, query_pb, namespace=None,

diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py
@@ -194,7 +194,8 @@ def test_lookup_single_key_empty_response(self):
             'lookup',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        self.assertEqual(conn.lookup(DATASET_ID, key_pb), None)
+        found = conn.lookup(DATASET_ID, [key_pb])
+        self.assertEqual(len(found), 0)
         cw = http._called_with
         self._verifyProtobufCall(cw, URI, conn)
         rq_class = datastore_pb.LookupRequest
@@ -220,7 +221,8 @@ def test_lookup_single_key_empty_response_w_eventual(self):
             'lookup',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        self.assertEqual(conn.lookup(DATASET_ID, key_pb, eventual=True), None)
+        found = conn.lookup(DATASET_ID, [key_pb], eventual=True)
+        self.assertEqual(len(found), 0)
         cw = http._called_with
         self._verifyProtobufCall(cw, URI, conn)
         rq_class = datastore_pb.LookupRequest
@@ -258,8 +260,8 @@ def test_lookup_single_key_empty_response_w_transaction(self):
             'lookup',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        found = conn.lookup(DATASET_ID, key_pb, transaction_id=TRANSACTION)
-        self.assertEqual(found, None)
+        found = conn.lookup(DATASET_ID, [key_pb], transaction_id=TRANSACTION)
+        self.assertEqual(len(found), 0)
         cw = http._called_with
         self._verifyProtobufCall(cw, URI, conn)
         rq_class = datastore_pb.LookupRequest
@@ -289,7 +291,7 @@ def test_lookup_single_key_nonempty_response(self):
             'lookup',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        found = conn.lookup(DATASET_ID, key_pb)
+        found, = conn.lookup(DATASET_ID, [key_pb])
         self.assertEqual(found.key.path_element[0].kind, 'Kind')
         self.assertEqual(found.key.path_element[0].id, 1234)
         cw = http._called_with