From 406ecad4719b3591427b1201ebb81b4d5968cd4e Mon Sep 17 00:00:00 2001 From: Jon Wayne Parrott Date: Fri, 18 Aug 2017 13:40:31 -0700 Subject: [PATCH] Add google.api.page_iterator.GRPCIterator (#3843) --- core/google/api/core/page_iterator.py | 87 +++++++++++++++++++ .../tests/unit/api_core/test_page_iterator.py | 84 ++++++++++++++++++ 2 files changed, 171 insertions(+) diff --git a/core/google/api/core/page_iterator.py b/core/google/api/core/page_iterator.py index 147c9f47e35a..23c469f9bc1d 100644 --- a/core/google/api/core/page_iterator.py +++ b/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 diff --git a/core/tests/unit/api_core/test_page_iterator.py b/core/tests/unit/api_core/test_page_iterator.py index 82466579e37b..541e60a61ffd 100644 --- a/core/tests/unit/api_core/test_page_iterator.py +++ b/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):