Skip to content

Commit

Permalink
feat(api_core): provide a 'raw_page' field for page_iterator.Page (#9486
Browse files Browse the repository at this point in the history
)

* Provide a 'raw_page' field for page_iterator.Page

Some paginated response messages include additional fields that users
may wish to inspect.
  • Loading branch information
software-dov authored and busunkim96 committed Nov 1, 2019
1 parent df185b9 commit 21e9457
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 10 deletions.
14 changes: 11 additions & 3 deletions api_core/google/api_core/page_iterator.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,14 +96,22 @@ class Page(object):
Callable to convert an item from the type in the raw API response
into the native object. Will be called with the iterator and a
single item.
raw_page Optional[google.protobuf.message.Message]:
The raw page response.
"""

def __init__(self, parent, items, item_to_value):
def __init__(self, parent, items, item_to_value, raw_page=None):
self._parent = parent
self._num_items = len(items)
self._remaining = self._num_items
self._item_iter = iter(items)
self._item_to_value = item_to_value
self._raw_page = raw_page

@property
def raw_page(self):
"""google.protobuf.message.Message"""
return self._raw_page

@property
def num_items(self):
Expand Down Expand Up @@ -360,7 +368,7 @@ def _next_page(self):
if self._has_next_page():
response = self._get_next_page_response()
items = response.get(self._items_key, ())
page = Page(self, items, self.item_to_value)
page = Page(self, items, self.item_to_value, raw_page=response)
self._page_start(self, page, response)
self.next_page_token = response.get(self._next_token)
return page
Expand Down Expand Up @@ -527,7 +535,7 @@ def _next_page(self):

self.next_page_token = getattr(response, self._response_token_field)
items = getattr(response, self._items_field)
page = Page(self, items, self.item_to_value)
page = Page(self, items, self.item_to_value, raw_page=response)

return page

Expand Down
24 changes: 17 additions & 7 deletions api_core/tests/unit/test_page_iterator.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,10 @@ def test_constructor(self):
assert page.remaining == 3
assert page._parent is parent
assert page._item_to_value is item_to_value
assert page.raw_page is None

def test___iter__(self):
page = page_iterator.Page(None, (), None)
page = page_iterator.Page(None, (), None, None)
assert iter(page) is page

def test_iterator_calls_parent_item_to_value(self):
Expand Down Expand Up @@ -69,6 +70,18 @@ def test_iterator_calls_parent_item_to_value(self):
item_to_value.assert_called_with(parent, 12)
assert page.remaining == 97

def test_raw_page(self):
parent = mock.sentinel.parent
item_to_value = mock.sentinel.item_to_value

raw_page = mock.sentinel.raw_page

page = page_iterator.Page(parent, (1, 2, 3), item_to_value, raw_page=raw_page)
assert page.raw_page is raw_page

with pytest.raises(AttributeError):
page.raw_page = None


class PageIteratorImpl(page_iterator.Iterator):
def _next_page(self):
Expand Down Expand Up @@ -116,8 +129,7 @@ def test_pages_property_restart(self):
def test__page_iter_increment(self):
iterator = PageIteratorImpl(None, None)
page = page_iterator.Page(
iterator, ("item",), page_iterator._item_to_value_identity
)
iterator, ("item",), page_iterator._item_to_value_identity)
iterator._next_page = mock.Mock(side_effect=[page, None])

assert iterator.num_results == 0
Expand Down Expand Up @@ -147,11 +159,9 @@ def test__items_iter(self):
# Make pages from mock responses
parent = mock.sentinel.parent
page1 = page_iterator.Page(
parent, (item1, item2), page_iterator._item_to_value_identity
)
parent, (item1, item2), page_iterator._item_to_value_identity)
page2 = page_iterator.Page(
parent, (item3,), page_iterator._item_to_value_identity
)
parent, (item3,), page_iterator._item_to_value_identity)

iterator = PageIteratorImpl(None, None)
iterator._next_page = mock.Mock(side_effect=[page1, page2, None])
Expand Down

0 comments on commit 21e9457

Please sign in to comment.