Add async iteration to v2 list auto-pagination 🪿✨

stripe/v2/_list_object.py

@@ -1,5 +1,6 @@
+from stripe._any_iterator import AnyIterator
 from stripe._stripe_object import StripeObject
-from typing import List, Optional, TypeVar, Generic
+from typing import AsyncIterator, Iterator, List, Optional, TypeVar, Generic
 
 
 T = TypeVar("T", bound=StripeObject)
@@ -40,7 +41,13 @@ def __len__(self):
     def __reversed__(self):
         return getattr(self, "data", []).__reversed__()
 
-    def auto_paging_iter(self):
+    def auto_paging_iter(self) -> AnyIterator[T]:
+        return AnyIterator(
+            self._auto_paging_iter(),
+            self._auto_paging_iter_async(),
+        )
+
+    def _auto_paging_iter(self) -> Iterator[T]:
         page = self.data
         next_page_url = self.next_page_url
         while True:
@@ -57,3 +64,21 @@ def auto_paging_iter(self):
             assert isinstance(result, ListObject)
             page = result.data
             next_page_url = result.next_page_url
+
+    async def _auto_paging_iter_async(self) -> AsyncIterator[T]:
+        page = self.data
+        next_page_url = self.next_page_url
+        while True:
+            for item in page:
+                yield item
+            if next_page_url is None:
+                break
+
+            result = await self._request_async(
+                "get",
+                next_page_url,
+                base_address="api",
+            )
+            assert isinstance(result, ListObject)
+            page = result.data
+            next_page_url = result.next_page_url

tests/api_resources/test_list_object_v2.py

@@ -146,3 +146,129 @@ def test_iter_forwards_api_key(self, http_client_mock: HTTPClientMock):
             query_string=query_string_2,
             api_key="sk_test_iter_forwards_options",
         )
+
+
+class TestAutoPagingAsync:
+    @staticmethod
+    def pageable_model_response(ids, next_page_url):
+        return {
+            "data": [{"id": id, "object": "pageablemodel"} for id in ids],
+            "next_page_url": next_page_url,
+        }
+
+    @pytest.mark.anyio
+    async def test_iter_one_page(self, http_client_mock):
+        lo = ListObject.construct_from(
+            self.pageable_model_response(["pm_123", "pm_124"], None), "mykey"
+        )
+
+        http_client_mock.assert_no_request()
+
+        seen = [item["id"] async for item in lo.auto_paging_iter()]
+
+        assert seen == ["pm_123", "pm_124"]
+
+    @pytest.mark.anyio
+    async def test_iter_two_pages(self, http_client_mock):
+        method = "get"
+        path = "/v2/pageablemodels"
+
+        lo = ListObject.construct_from(
+            self.pageable_model_response(
+                ["pm_123", "pm_124"],
+                "/v2/pageablemodels?foo=bar&page=page_2",
+            ),
+            None,
+        )
+
+        http_client_mock.stub_request(
+            method,
+            path=path,
+            query_string="foo=bar&page=page_3",
+            rbody=json.dumps(
+                self.pageable_model_response(["pm_127", "pm_128"], None)
+            ),
+        )
+
+        http_client_mock.stub_request(
+            method,
+            path=path,
+            query_string="foo=bar&page=page_2",
+            rbody=json.dumps(
+                self.pageable_model_response(
+                    ["pm_125", "pm_126"],
+                    "/v2/pageablemodels?foo=bar&page=page_3",
+                )
+            ),
+        )
+
+        seen = [item["id"] async for item in lo.auto_paging_iter()]
+
+        http_client_mock.assert_requested(
+            method, path=path, query_string="foo=bar&page=page_2"
+        )
+        http_client_mock.assert_requested(
+            method, path=path, query_string="foo=bar&page=page_3"
+        )
+
+        assert seen == [
+            "pm_123",
+            "pm_124",
+            "pm_125",
+            "pm_126",
+            "pm_127",
+            "pm_128",
+        ]
+
+    @pytest.mark.anyio
+    async def test_iter_forwards_api_key(
+        self, http_client_mock: HTTPClientMock
+    ):
+        client = stripe.StripeClient(
+            http_client=http_client_mock.get_mock_http_client(),
+            api_key="sk_test_xyz",
+        )
+
+        method = "get"
+        query_string_1 = "object_id=obj_123"
+        query_string_2 = "object_id=obj_123&page=page_2"
+        path = "/v2/core/events"
+
+        http_client_mock.stub_request(
+            method,
+            path=path,
+            query_string=query_string_1,
+            rbody='{"data": [{"id": "x"}], "next_page_url": "/v2/core/events?object_id=obj_123&page=page_2"}',
+            rcode=200,
+            rheaders={},
+        )
+
+        http_client_mock.stub_request(
+            method,
+            path=path,
+            query_string=query_string_2,
+            rbody='{"data": [{"id": "y"}, {"id": "z"}], "next_page_url": null}',
+            rcode=200,
+            rheaders={},
+        )
+
+        lo = await client.v2.core.events.list_async(
+            params={"object_id": "obj_123"},
+            options={"api_key": "sk_test_iter_forwards_options"},
+        )
+
+        seen = [item["id"] async for item in lo.auto_paging_iter()]
+
+        assert seen == ["x", "y", "z"]
+        http_client_mock.assert_requested(
+            method,
+            path=path,
+            query_string=query_string_1,
+            api_key="sk_test_iter_forwards_options",
+        )
+        http_client_mock.assert_requested(
+            method,
+            path=path,
+            query_string=query_string_2,
+            api_key="sk_test_iter_forwards_options",
+        )