diff --git a/api_core/google/api_core/page_iterator.py b/api_core/google/api_core/page_iterator.py index 3ac5904399b0..11a92d38f3ce 100644 --- a/api_core/google/api_core/page_iterator.py +++ b/api_core/google/api_core/page_iterator.py @@ -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): @@ -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 @@ -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 diff --git a/api_core/tests/unit/test_page_iterator.py b/api_core/tests/unit/test_page_iterator.py index 6335001bcf41..2bf742492889 100644 --- a/api_core/tests/unit/test_page_iterator.py +++ b/api_core/tests/unit/test_page_iterator.py @@ -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): @@ -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): @@ -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 @@ -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])