11from __future__ import annotations
22
3- from collections .abc import AsyncIterable , AsyncIterator , Awaitable , Generator
3+ from collections .abc import AsyncIterable , AsyncIterator , Awaitable , Generator , Iterable , Iterator
44from typing import Any , Generic , Protocol , TypeVar
55
66from apify_client ._utils import parse_date_fields , pluck_data
@@ -52,6 +52,37 @@ def _list(self, **kwargs: Any) -> ListPage:
5252
5353 return ListPage (parse_date_fields (pluck_data (response .json ())))
5454
55+
56+ def _list_iterable (self , ** kwargs : Any ) -> IterableListPage [T ]:
57+ """Return object can be awaited or iterated over."""
58+ chunk_size = kwargs .pop ('chunk_size' , None )
59+
60+ list_page = self ._list (** {** kwargs , 'limit' : _min_for_limit_param (kwargs .get ('limit' ), chunk_size )})
61+
62+ def iterator () -> Iterator [T ]:
63+ current_page = list_page
64+ for item in current_page .items :
65+ yield item
66+
67+ offset = kwargs .get ('offset' ) or 0
68+ limit = min (kwargs .get ('limit' ) or current_page .total , current_page .total )
69+
70+ current_offset = offset + len (current_page .items )
71+ remaining_items = min (current_page .total - offset , limit ) - len (current_page .items )
72+ while current_page .items and remaining_items > 0 :
73+ new_kwargs = {
74+ ** kwargs ,
75+ 'offset' : current_offset ,
76+ 'limit' : _min_for_limit_param (remaining_items , chunk_size ),
77+ }
78+ current_page = self ._list (** new_kwargs )
79+ for item in current_page .items :
80+ yield item
81+ current_offset += len (current_page .items )
82+ remaining_items -= len (current_page .items )
83+
84+ return IterableListPage [T ](list_page , iterator ())
85+
5586 def _create (self , resource : dict ) -> dict :
5687 response = self .http_client .call (
5788 url = self ._url (),
@@ -85,24 +116,11 @@ async def _list(self, **kwargs: Any) -> ListPage:
85116
86117 return ListPage (parse_date_fields (pluck_data (response .json ())))
87118
88- def _list_iterable (self , ** kwargs : Any ) -> ListPageProtocol [T ]:
119+ def _list_iterable (self , ** kwargs : Any ) -> ListPageProtocolAsync [T ]:
89120 """Return object can be awaited or iterated over."""
90-
91- def min_for_limit_param (a : int | None , b : int | None ) -> int | None :
92- # API treats 0 as None for limit parameter, in this context API understands 0 as infinity.
93- if a == 0 :
94- a = None
95- if b == 0 :
96- b = None
97- if a is None :
98- return b
99- if b is None :
100- return a
101- return min (a , b )
102-
103121 chunk_size = kwargs .pop ('chunk_size' , None )
104122
105- list_page_awaitable = self ._list (** {** kwargs , 'limit' : min_for_limit_param (kwargs .get ('limit' ), chunk_size )})
123+ list_page_awaitable = self ._list (** {** kwargs , 'limit' : _min_for_limit_param (kwargs .get ('limit' ), chunk_size )})
106124
107125 async def async_iterator () -> AsyncIterator [T ]:
108126 current_page = await list_page_awaitable
@@ -118,15 +136,15 @@ async def async_iterator() -> AsyncIterator[T]:
118136 new_kwargs = {
119137 ** kwargs ,
120138 'offset' : current_offset ,
121- 'limit' : min_for_limit_param (remaining_items , chunk_size ),
139+ 'limit' : _min_for_limit_param (remaining_items , chunk_size ),
122140 }
123141 current_page = await self ._list (** new_kwargs )
124142 for item in current_page .items :
125143 yield item
126144 current_offset += len (current_page .items )
127145 remaining_items -= len (current_page .items )
128146
129- return IterableListPage [T ](list_page_awaitable , async_iterator ())
147+ return IterableListPageAsync [T ](list_page_awaitable , async_iterator ())
130148
131149 async def _create (self , resource : dict ) -> dict :
132150 response = await self .http_client .call (
@@ -153,11 +171,47 @@ async def _get_or_create(
153171 return parse_date_fields (pluck_data (response .json ()))
154172
155173
156- class ListPageProtocol (Protocol [T ], AsyncIterable [ T ], Awaitable [ ListPage [ T ] ]):
174+ class ListPageProtocol (Protocol [T ], Iterable [ T ]):
157175 """Protocol for an object that can be both awaited and asynchronously iterated over."""
158176
177+ items : list [T ]
178+ """List of returned objects on this page"""
159179
160- class IterableListPage (Generic [T ]):
180+ count : int
181+ """Count of the returned objects on this page"""
182+
183+ offset : int
184+ """The limit on the number of returned objects offset specified in the API call"""
185+
186+ limit : int
187+ """The offset of the first object specified in the API call"""
188+
189+ total : int
190+ """Total number of objects matching the API call criteria"""
191+
192+ desc : bool
193+ """Whether the listing is descending or not"""
194+
195+ class IterableListPage (Generic [T ], ListPage [T ]):
196+ """Can be called to get ListPage with items or iterated over to get individual items."""
197+
198+ def __init__ (self , list_page : ListPage [T ], iterator : Iterator [T ]) -> None :
199+ self .items = list_page .items
200+ self .offset = list_page .offset
201+ self .limit = list_page .limit
202+ self .count = list_page .count
203+ self .total = list_page .total
204+ self .desc = list_page .desc
205+ self ._iterator = iterator
206+
207+ def __iter__ (self ) -> Iterator [T ]:
208+ """Return an iterator over the items from API, possibly doing multiple API calls."""
209+ return self ._iterator
210+
211+ class ListPageProtocolAsync (Protocol [T ], AsyncIterable [T ], Awaitable [ListPage [T ]]):
212+ """Protocol for an object that can be both awaited and asynchronously iterated over."""
213+
214+ class IterableListPageAsync (Generic [T ]):
161215 """Can be awaited to get ListPage with items or asynchronously iterated over to get individual items."""
162216
163217 def __init__ (self , awaitable : Awaitable [ListPage [T ]], async_iterator : AsyncIterator [T ]) -> None :
@@ -171,3 +225,17 @@ def __aiter__(self) -> AsyncIterator[T]:
171225 def __await__ (self ) -> Generator [Any , Any , ListPage [T ]]:
172226 """Return an awaitable that resolves to the ListPage doing exactly one API call."""
173227 return self ._awaitable .__await__ ()
228+
229+
230+ def _min_for_limit_param (a : int | None , b : int | None ) -> int | None :
231+ """Return minimum of two limit parameters, treating None or 0 as infinity. Return None for infinity."""
232+ # API treats 0 as None for limit parameter, in this context API understands 0 as infinity.
233+ if a == 0 :
234+ a = None
235+ if b == 0 :
236+ b = None
237+ if a is None :
238+ return b
239+ if b is None :
240+ return a
241+ return min (a , b )
0 commit comments