Made kind take and return a single string and change more_results.

Incorporates feedback from @pcostell. In particular, treats
`NOT_FINISHED` as the only enum that needs `more_results=True`
and both `NO_MORE_RESULTS` and `MORE_RESULTS_AFTER_LIMIT` as
responses which don't need paging.

Author: dhermes

gcloud/datastore/query.py

@@ -56,8 +56,11 @@ class Query(object):
     :param dataset: The namespace to which to restrict results.
     """
 
-    _MORE_RESULTS = datastore_pb.QueryResultBatch.MORE_RESULTS_AFTER_LIMIT
-    _NO_MORE_RESULTS = datastore_pb.QueryResultBatch.NO_MORE_RESULTS
+    _NOT_FINISHED = datastore_pb.QueryResultBatch.NOT_FINISHED
+    _FINISHED = (
+        datastore_pb.QueryResultBatch.NO_MORE_RESULTS,
+        datastore_pb.QueryResultBatch.MORE_RESULTS_AFTER_LIMIT,
+    )
     OPERATORS = {
         '<=': datastore_pb.PropertyFilter.LESS_THAN_OR_EQUAL,
         '>=': datastore_pb.PropertyFilter.GREATER_THAN_OR_EQUAL,
@@ -227,26 +230,24 @@ def ancestor(self, ancestor):
 
         return clone
 
-    def kind(self, *kinds):
+    def kind(self, kind=None):
         """Get or set the Kind of the Query.
 
-        .. note::
-          This is an **additive** operation.  That is, if the Query is
-          set for kinds A and B, and you call ``.kind('C')``, it will
-          query for kinds A, B, *and*, C.
+        :type kind: string
+        :param kind: Optional. The entity kinds for which to query.
 
-        :type kinds: string
-        :param kinds: The entity kinds for which to query.
-
-        :rtype: string, list of strings, or :class:`Query`
-        :returns: If no arguments, returns the kind or list of kinds.
-                  If a kind is provided, returns a clone of the :class:`Query`
-                  with those kinds set.
+        :rtype: string or :class:`Query`
+        :returns: If `kind` is None, returns the kind. If a kind is provided,
+                  returns a clone of the :class:`Query` with that kind set.
+        :raises: `ValueError` from the getter if multiple kinds are set on
+                 the query.
         """
-        if kinds:
+        if kind is not None:
+            kinds = [kind]
             clone = self._clone()
-            for kind in kinds:
-                clone._pb.kind.add().name = kind
+            clone._pb.ClearField('kind')
+            for new_kind in kinds:
+                clone._pb.kind.add().name = new_kind
             return clone
         else:
             # In the proto definition for Query, `kind` is repeated.
@@ -255,7 +256,7 @@ def kind(self, *kinds):
             if num_kinds == 1:
                 return kind_names[0]
             elif num_kinds > 1:
-                return kind_names
+                raise ValueError('Only a single kind can be set.')
 
     def limit(self, limit=None):
         """Get or set the limit of the Query.
@@ -343,7 +344,7 @@ def fetch_page(self, limit=None):
                   encoded cursor for paging and the third is a boolean
                   indicating if there are more results.
         :raises: `ValueError` if more_results is not one of the enums
-                 MORE_RESULTS_AFTER_LIMIT or NO_MORE_RESULTS.
+                 NOT_FINISHED, MORE_RESULTS_AFTER_LIMIT, NO_MORE_RESULTS.
         """
         clone = self
 
@@ -372,13 +373,11 @@ def fetch_page(self, limit=None):
 
         cursor = base64.b64encode(cursor_as_bytes)
 
-        if more_results_enum == self._MORE_RESULTS:
+        if more_results_enum == self._NOT_FINISHED:
             more_results = True
-        elif more_results_enum == self._NO_MORE_RESULTS:
+        elif more_results_enum in self._FINISHED:
             more_results = False
         else:
-            # Note this covers the value NOT_FINISHED since this fetch does
-            # not occur within a batch, we don't expect to see NOT_FINISHED.
             raise ValueError('Unexpected value returned for `more_results`.')
 
         return entities, cursor, more_results

gcloud/datastore/test_query.py

@@ -261,11 +261,22 @@ def test_kind_setter_w_existing(self):
         _KIND_AFTER = 'KIND_AFTER'
         dataset = Dataset(_DATASET)
         query = self._makeOne(_KIND_BEFORE, dataset)
+        self.assertEqual(query.kind(), _KIND_BEFORE)
         after = query.kind(_KIND_AFTER)
         self.assertFalse(after is query)
         self.assertTrue(isinstance(after, self._getTargetClass()))
         self.assertTrue(after.dataset() is dataset)
-        self.assertEqual(after.kind(), [_KIND_BEFORE, _KIND_AFTER])
+        self.assertEqual(after.kind(), _KIND_AFTER)
+
+    def test_kind_getter_unset(self):
+        query = self._makeOne()
+        self.assertEqual(query.kind(), None)
+
+    def test_kind_getter_bad_pb(self):
+        query = self._makeOne()
+        query._pb.kind.add().name = 'foo'
+        query._pb.kind.add().name = 'bar'
+        self.assertRaises(ValueError, query.kind)
 
     def test_limit_setter_wo_existing(self):
         from gcloud.datastore.dataset import Dataset
@@ -293,8 +304,9 @@ def test_dataset_setter(self):
         self.assertTrue(after.dataset() is dataset)
         self.assertEqual(query.kind(), _KIND)
 
-    def _fetch_page_helper(self, cursor=b'\x00', limit=None, more_results=True,
-                           _more_pb=None, use_fetch=False):
+    def _fetch_page_helper(self, cursor=b'\x00', limit=None,
+                           more_results=False, _more_pb=None,
+                           use_fetch=False):
         import base64
         from gcloud.datastore.datastore_v1_pb2 import Entity
         _CURSOR_FOR_USER = (None if cursor is None
@@ -351,16 +363,14 @@ def test_fetch_page_explicit_limit(self):
     def test_fetch_page_no_more_results(self):
         from gcloud.datastore import datastore_v1_pb2 as datastore_pb
         no_more = datastore_pb.QueryResultBatch.NO_MORE_RESULTS
-        self._fetch_page_helper(cursor='CURSOR', limit=13, more_results=False,
-                                _more_pb=no_more)
+        self._fetch_page_helper(cursor='CURSOR', limit=13, _more_pb=no_more)
 
-    def test_fetch_page_more_results_ill_formed(self):
+    def test_fetch_page_not_finished(self):
         from gcloud.datastore import datastore_v1_pb2 as datastore_pb
         not_finished = datastore_pb.QueryResultBatch.NOT_FINISHED
-        # Try a valid enum but not allowed.
-        self.assertRaises(ValueError, self._fetch_page_helper,
-                          _more_pb=not_finished)
-        # Try an invalid enum but not allowed.
+        self._fetch_page_helper(_more_pb=not_finished, more_results=True)
+
+    def test_fetch_page_more_results_invalid(self):
         self.assertRaises(ValueError, self._fetch_page_helper,
                           _more_pb=object())