feat: add iterator capability to paged iterators (#200)

plamut · web-flow · commit 3487d68bdab6 · 2021-06-08T10:30:29.000+02:00
* 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
diff --git a/google/api_core/page_iterator.py b/google/api_core/page_iterator.py
@@ -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
@@ -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.
 
diff --git a/google/api_core/page_iterator_async.py b/google/api_core/page_iterator_async.py
@@ -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
@@ -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.
 
diff --git a/tests/asyncio/test_page_iterator_async.py b/tests/asyncio/test_page_iterator_async.py
@@ -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)
 
diff --git a/tests/unit/test_page_iterator.py b/tests/unit/test_page_iterator.py
@@ -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)
 
