Skip to content

Commit

Permalink
feat: add iterator capability to paged iterators (#200)
Browse files Browse the repository at this point in the history
* feat: add iterator capability to *Iterator classes

The *Iterator classes are only _iterables_, and this commit also makes
them _iterators_, i.e. calling next(iterator) on them now works.

* Make AsyncIterator an actual async iterator
  • Loading branch information
plamut committed Jun 8, 2021
1 parent 641fbbf commit 3487d68
Show file tree
Hide file tree
Showing 4 changed files with 64 additions and 0 deletions.
10 changes: 10 additions & 0 deletions google/api_core/page_iterator.py
Expand Up @@ -170,6 +170,8 @@ def __init__(
max_results=None,
):
self._started = False
self.__active_iterator = None

self.client = client
"""Optional[Any]: The client that created this iterator."""
self.item_to_value = item_to_value
Expand Down Expand Up @@ -228,6 +230,14 @@ def __iter__(self):
self._started = True
return self._items_iter()

def __next__(self):
if self.__active_iterator is None:
self.__active_iterator = iter(self)
return next(self.__active_iterator)

# Preserve Python 2 compatibility.
next = __next__

def _page_iter(self, increment):
"""Generator of pages of API responses.
Expand Down
7 changes: 7 additions & 0 deletions google/api_core/page_iterator_async.py
Expand Up @@ -101,6 +101,8 @@ def __init__(
max_results=None,
):
self._started = False
self.__active_aiterator = None

self.client = client
"""Optional[Any]: The client that created this iterator."""
self.item_to_value = item_to_value
Expand Down Expand Up @@ -159,6 +161,11 @@ def __aiter__(self):
self._started = True
return self._items_aiter()

async def __anext__(self):
if self.__active_aiterator is None:
self.__active_aiterator = self.__aiter__()
return await self.__active_aiterator.__anext__()

async def _page_aiter(self, increment):
"""Generator of pages of API responses.
Expand Down
27 changes: 27 additions & 0 deletions tests/asyncio/test_page_iterator_async.py
Expand Up @@ -47,6 +47,33 @@ def test_constructor(self):
assert iterator.next_page_token == token
assert iterator.num_results == 0

@pytest.mark.asyncio
async def test_anext(self):
parent = mock.sentinel.parent
page_1 = page_iterator_async.Page(
parent, ("item 1.1", "item 1.2"), page_iterator_async._item_to_value_identity
)
page_2 = page_iterator_async.Page(
parent, ("item 2.1",), page_iterator_async._item_to_value_identity
)

async_iterator = PageAsyncIteratorImpl(None, None)
async_iterator._next_page = mock.AsyncMock(side_effect=[page_1, page_2, None])

# Consume items and check the state of the async_iterator.
assert async_iterator.num_results == 0
assert await async_iterator.__anext__() == "item 1.1"
assert async_iterator.num_results == 1

assert await async_iterator.__anext__() == "item 1.2"
assert async_iterator.num_results == 2

assert await async_iterator.__anext__() == "item 2.1"
assert async_iterator.num_results == 3

with pytest.raises(StopAsyncIteration):
await async_iterator.__anext__()

def test_pages_property_starts(self):
iterator = PageAsyncIteratorImpl(None, None)

Expand Down
20 changes: 20 additions & 0 deletions tests/unit/test_page_iterator.py
Expand Up @@ -109,6 +109,26 @@ def test_constructor(self):
assert iterator.next_page_token == token
assert iterator.num_results == 0

def test_next(self):
iterator = PageIteratorImpl(None, None)
page_1 = page_iterator.Page(
iterator, ("item 1.1", "item 1.2"), page_iterator._item_to_value_identity
)
page_2 = page_iterator.Page(
iterator, ("item 2.1",), page_iterator._item_to_value_identity
)
iterator._next_page = mock.Mock(side_effect=[page_1, page_2, None])

result = next(iterator)
assert result == "item 1.1"
result = next(iterator)
assert result == "item 1.2"
result = next(iterator)
assert result == "item 2.1"

with pytest.raises(StopIteration):
next(iterator)

def test_pages_property_starts(self):
iterator = PageIteratorImpl(None, None)

Expand Down

0 comments on commit 3487d68

Please sign in to comment.