Skip to content

Commit 9230dd4

Browse files
software-dovbusunkim96
authored andcommitted
feat(api_core): provide a 'raw_page' field for page_iterator.Page (#9486)
* Provide a 'raw_page' field for page_iterator.Page Some paginated response messages include additional fields that users may wish to inspect.
1 parent 6ac4dc2 commit 9230dd4

File tree

2 files changed

+28
-10
lines changed

2 files changed

+28
-10
lines changed

google/api_core/page_iterator.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -96,14 +96,22 @@ class Page(object):
9696
Callable to convert an item from the type in the raw API response
9797
into the native object. Will be called with the iterator and a
9898
single item.
99+
raw_page Optional[google.protobuf.message.Message]:
100+
The raw page response.
99101
"""
100102

101-
def __init__(self, parent, items, item_to_value):
103+
def __init__(self, parent, items, item_to_value, raw_page=None):
102104
self._parent = parent
103105
self._num_items = len(items)
104106
self._remaining = self._num_items
105107
self._item_iter = iter(items)
106108
self._item_to_value = item_to_value
109+
self._raw_page = raw_page
110+
111+
@property
112+
def raw_page(self):
113+
"""google.protobuf.message.Message"""
114+
return self._raw_page
107115

108116
@property
109117
def num_items(self):
@@ -360,7 +368,7 @@ def _next_page(self):
360368
if self._has_next_page():
361369
response = self._get_next_page_response()
362370
items = response.get(self._items_key, ())
363-
page = Page(self, items, self.item_to_value)
371+
page = Page(self, items, self.item_to_value, raw_page=response)
364372
self._page_start(self, page, response)
365373
self.next_page_token = response.get(self._next_token)
366374
return page
@@ -527,7 +535,7 @@ def _next_page(self):
527535

528536
self.next_page_token = getattr(response, self._response_token_field)
529537
items = getattr(response, self._items_field)
530-
page = Page(self, items, self.item_to_value)
538+
page = Page(self, items, self.item_to_value, raw_page=response)
531539

532540
return page
533541

tests/unit/test_page_iterator.py

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,10 @@ def test_constructor(self):
3636
assert page.remaining == 3
3737
assert page._parent is parent
3838
assert page._item_to_value is item_to_value
39+
assert page.raw_page is None
3940

4041
def test___iter__(self):
41-
page = page_iterator.Page(None, (), None)
42+
page = page_iterator.Page(None, (), None, None)
4243
assert iter(page) is page
4344

4445
def test_iterator_calls_parent_item_to_value(self):
@@ -69,6 +70,18 @@ def test_iterator_calls_parent_item_to_value(self):
6970
item_to_value.assert_called_with(parent, 12)
7071
assert page.remaining == 97
7172

73+
def test_raw_page(self):
74+
parent = mock.sentinel.parent
75+
item_to_value = mock.sentinel.item_to_value
76+
77+
raw_page = mock.sentinel.raw_page
78+
79+
page = page_iterator.Page(parent, (1, 2, 3), item_to_value, raw_page=raw_page)
80+
assert page.raw_page is raw_page
81+
82+
with pytest.raises(AttributeError):
83+
page.raw_page = None
84+
7285

7386
class PageIteratorImpl(page_iterator.Iterator):
7487
def _next_page(self):
@@ -116,8 +129,7 @@ def test_pages_property_restart(self):
116129
def test__page_iter_increment(self):
117130
iterator = PageIteratorImpl(None, None)
118131
page = page_iterator.Page(
119-
iterator, ("item",), page_iterator._item_to_value_identity
120-
)
132+
iterator, ("item",), page_iterator._item_to_value_identity)
121133
iterator._next_page = mock.Mock(side_effect=[page, None])
122134

123135
assert iterator.num_results == 0
@@ -147,11 +159,9 @@ def test__items_iter(self):
147159
# Make pages from mock responses
148160
parent = mock.sentinel.parent
149161
page1 = page_iterator.Page(
150-
parent, (item1, item2), page_iterator._item_to_value_identity
151-
)
162+
parent, (item1, item2), page_iterator._item_to_value_identity)
152163
page2 = page_iterator.Page(
153-
parent, (item3,), page_iterator._item_to_value_identity
154-
)
164+
parent, (item3,), page_iterator._item_to_value_identity)
155165

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

0 commit comments

Comments
 (0)