Draft for testing

Author: Pijukatel

src/apify_client/clients/base/resource_collection_client.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import Any, Generic, TypeVar
+from collections.abc import AsyncIterator, Awaitable
+from typing import Any, Generic, Protocol, TypeVar
 
 from apify_client._utils import parse_date_fields, pluck_data
 from apify_client.clients.base.base_client import BaseClient, BaseClientAsync
@@ -84,6 +85,45 @@ async def _list(self, **kwargs: Any) -> ListPage:
 
         return ListPage(parse_date_fields(pluck_data(response.json())))
 
+    def _list_paginated(self, **kwargs: Any) -> ListPageProtocol:
+        def min_for_limit_param(a: int | None, b : int| None) -> int | None:
+            # API treats 0 as None for limit parameter, in this context API understands 0 as infinity.
+            if a == 0:
+                a = None
+            if b == 0:
+                b = None
+            if a is None:
+                return b
+            if b is None:
+                return a
+            return min(a, b)
+        chunk_size = kwargs.pop('chunk_size', None)
+
+        list_page_getter = self._list(**{**kwargs, 'limit':min_for_limit_param(kwargs.get('limit'), chunk_size)})
+
+        async def async_iterator():
+            current_page = await list_page_getter
+            for item in current_page.items:
+                yield item
+
+            offset = kwargs.get('offset') or 0
+            limit = min(kwargs.get('limit') or current_page.total, current_page.total)
+
+            current_offset = offset + len(current_page.items)
+            remaining_items = min(current_page.total-offset, limit) - len(current_page.items)
+            while (current_page.items and remaining_items > 0):
+                new_kwargs = {**kwargs,
+                              'offset': current_offset,
+                              'limit': min_for_limit_param(remaining_items, chunk_size)}
+                current_page = await self._list(**new_kwargs)
+                for item in current_page.items:
+                    yield item
+                current_offset += len(current_page.items)
+                remaining_items -= len(current_page.items)
+
+        return ListPageIterable(list_page_getter, async_iterator())
+
+
     async def _create(self, resource: dict) -> dict:
         response = await self.http_client.call(
             url=self._url(),
@@ -107,3 +147,27 @@ async def _get_or_create(
         )
 
         return parse_date_fields(pluck_data(response.json()))
+
+
+class ListPageProtocol(Protocol[T]):
+    def __aiter__(self) -> AsyncIterator[T]: ...
+    def __await__(self) -> ListPage[T]: ...
+
+
+class ListPageIterable(Generic[T]):
+    def __init__(self, awaitable: Awaitable[ListPage[T]], async_iterator: AsyncIterator[T]) -> None:
+        self._awaitable = awaitable
+        self._async_iterator = async_iterator
+
+    def  __aiter__(self):
+        return self._async_iterator
+
+    def __await__(self):
+        return self._awaitable.__await__()
+
+
+    """
+    async def __anext__(self) -> T:
+        async for item in self._async_iterator:
+            print(item)
+    """

src/apify_client/clients/resource_clients/actor_collection.py

@@ -7,7 +7,7 @@
 from apify_client.clients.resource_clients.actor import get_actor_representation
 
 if TYPE_CHECKING:
-    from apify_client.clients.base.resource_collection_client import ListPage
+    from apify_client.clients.base.resource_collection_client import ListPage, ListPageProtocol
 
 
 class ActorCollectionClient(ResourceCollectionClient):
@@ -142,15 +142,16 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         resource_path = kwargs.pop('resource_path', 'acts')
         super().__init__(*args, resource_path=resource_path, **kwargs)
 
-    async def list(
+    def list(
         self,
         *,
         my: bool | None = None,
         limit: int | None = None,
         offset: int | None = None,
         desc: bool | None = None,
         sort_by: Literal['createdAt', 'stats.lastRunStartedAt'] | None = 'createdAt',
-    ) -> ListPage[dict]:
+        chunk_size: int | None = None,
+    ) -> ListPageProtocol[dict]:
         """List the Actors the user has created or used.
 
         https://docs.apify.com/api/v2#/reference/actors/actor-collection/get-list-of-actors
@@ -165,7 +166,7 @@ async def list(
         Returns:
             The list of available Actors matching the specified filters.
         """
-        return await self._list(my=my, limit=limit, offset=offset, desc=desc, sortBy=sort_by)
+        return self._list_paginated(my=my, limit=limit, offset=offset, desc=desc, sortBy=sort_by, chunk_size=chunk_size)
 
     async def create(
         self,