Add google.api.page_iterator.GRPCIterator (googleapis#3843)

core/google/api/core/page_iterator.py

@@ -423,3 +423,90 @@ def _next_page(self):
             return page
         except StopIteration:
             return None
+
+
+class GRPCIterator(Iterator):
+    """A generic class for iterating through gRPC list responses.
+
+    .. note:: The class does not take a ``page_token`` argument because it can
+        just be specified in the ``request``.
+
+    Args:
+        client (google.cloud.client.Client): The API client. This unused by
+            this class, but kept to satisfy the :class:`Iterator` interface.
+        method (Callable[protobuf.Message]): A bound gRPC method that should
+            take a single message for the request.
+        request (protobuf.Message): The request message.
+        items_field (str): The field in the response message that has the
+            items for the page.
+        item_to_value (Callable[Iterator, Any]): Callable to convert an item
+            from the type in the JSON response into a native object. Will
+            be called with the iterator and a single item.
+        request_token_field (str): The field in the request message used to
+            specify the page token.
+        response_token_field (str): The field in the response message that has
+            the token for the next page.
+        max_results (int): The maximum number of results to fetch.
+
+    .. autoattribute:: pages
+    """
+
+    _DEFAULT_REQUEST_TOKEN_FIELD = 'page_token'
+    _DEFAULT_RESPONSE_TOKEN_FIELD = 'next_page_token'
+
+    def __init__(
+            self,
+            client,
+            method,
+            request,
+            items_field,
+            item_to_value=_item_to_value_identity,
+            request_token_field=_DEFAULT_REQUEST_TOKEN_FIELD,
+            response_token_field=_DEFAULT_RESPONSE_TOKEN_FIELD,
+            max_results=None):
+        super(GRPCIterator, self).__init__(
+            client, item_to_value, max_results=max_results)
+        self._method = method
+        self._request = request
+        self._items_field = items_field
+        self._request_token_field = request_token_field
+        self._response_token_field = response_token_field
+
+    def _next_page(self):
+        """Get the next page in the iterator.
+
+        Returns:
+            Page: The next page in the iterator or :data:`None` if there are no
+                pages left.
+        """
+        if not self._has_next_page():
+            return None
+
+        if self.next_page_token is not None:
+            setattr(
+                self._request, self._request_token_field, self.next_page_token)
+
+        response = self._method(self._request)
+
+        self.next_page_token = getattr(response, self._response_token_field)
+        items = getattr(response, self._items_field)
+        page = Page(self, items, self._item_to_value)
+
+        return page
+
+    def _has_next_page(self):
+        """Determines whether or not there are more pages with results.
+
+        Returns:
+            bool: Whether the iterator has more pages.
+        """
+        if self.page_number == 0:
+            return True
+
+        if self.max_results is not None:
+            if self.num_results >= self.max_results:
+                return False
+
+        # Note: intentionally a falsy check instead of a None check. The RPC
+        # can return an empty string indicating no more pages.
+        return True if self.next_page_token else False

core/tests/unit/api_core/test_page_iterator.py

@@ -408,6 +408,90 @@ def test__get_next_page_bad_http_method(self):
             iterator._get_next_page_response()
 
 
+class TestGRPCIterator(object):
+
+    def test_constructor(self):
+        client = mock.sentinel.client
+        items_field = 'items'
+        iterator = page_iterator.GRPCIterator(
+            client, mock.sentinel.method, mock.sentinel.request, items_field)
+
+        assert not iterator._started
+        assert iterator.client is client
+        assert iterator.max_results is None
+        assert iterator._method == mock.sentinel.method
+        assert iterator._request == mock.sentinel.request
+        assert iterator._items_field == items_field
+        assert iterator._item_to_value is page_iterator._item_to_value_identity
+        assert (iterator._request_token_field ==
+                page_iterator.GRPCIterator._DEFAULT_REQUEST_TOKEN_FIELD)
+        assert (iterator._response_token_field ==
+                page_iterator.GRPCIterator._DEFAULT_RESPONSE_TOKEN_FIELD)
+        # Changing attributes.
+        assert iterator.page_number == 0
+        assert iterator.next_page_token is None
+        assert iterator.num_results == 0
+
+    def test_constructor_options(self):
+        client = mock.sentinel.client
+        items_field = 'items'
+        request_field = 'request'
+        response_field = 'response'
+        iterator = page_iterator.GRPCIterator(
+            client, mock.sentinel.method, mock.sentinel.request, items_field,
+            item_to_value=mock.sentinel.item_to_value,
+            request_token_field=request_field,
+            response_token_field=response_field,
+            max_results=42)
+
+        assert iterator.client is client
+        assert iterator.max_results == 42
+        assert iterator._method == mock.sentinel.method
+        assert iterator._request == mock.sentinel.request
+        assert iterator._items_field == items_field
+        assert iterator._item_to_value is mock.sentinel.item_to_value
+        assert iterator._request_token_field == request_field
+        assert iterator._response_token_field == response_field
+
+    def test_iterate(self):
+        request = mock.Mock(spec=['page_token'], page_token=None)
+        response1 = mock.Mock(items=['a', 'b'], next_page_token='1')
+        response2 = mock.Mock(items=['c'], next_page_token='2')
+        response3 = mock.Mock(items=['d'], next_page_token='')
+        method = mock.Mock(side_effect=[response1, response2, response3])
+        iterator = page_iterator.GRPCIterator(
+            mock.sentinel.client, method, request, 'items')
+
+        assert iterator.num_results == 0
+
+        items = list(iterator)
+        assert items == ['a', 'b', 'c', 'd']
+
+        method.assert_called_with(request)
+        assert method.call_count == 3
+        assert request.page_token == '2'
+
+    def test_iterate_with_max_results(self):
+        request = mock.Mock(spec=['page_token'], page_token=None)
+        response1 = mock.Mock(items=['a', 'b'], next_page_token='1')
+        response2 = mock.Mock(items=['c'], next_page_token='2')
+        response3 = mock.Mock(items=['d'], next_page_token='')
+        method = mock.Mock(side_effect=[response1, response2, response3])
+        iterator = page_iterator.GRPCIterator(
+            mock.sentinel.client, method, request, 'items', max_results=3)
+
+        assert iterator.num_results == 0
+
+        items = list(iterator)
+
+        assert items == ['a', 'b', 'c']
+        assert iterator.num_results == 3
+
+        method.assert_called_with(request)
+        assert method.call_count == 2
+        assert request.page_token is '1'
+
+
 class GAXPageIterator(object):
     """Fake object that matches gax.PageIterator"""
     def __init__(self, pages, page_token=None):