# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Pagers for the GenAI List APIs."""
# pylint: disable=protected-access
import copy
from typing import Any, AsyncIterator, Awaitable, Callable, Generic, Iterator, Literal, TypeVar
T = TypeVar('T')
PagedItem = Literal[
'batch_jobs', 'models', 'tuning_jobs', 'files', 'cached_contents'
]
class _BasePager(Generic[T]):
"""Base pager class for iterating through paginated results."""
def __init__(
self,
name: PagedItem,
request: Callable[Any, Any],
response: Any,
config: Any,
):
self._name = name
self._request = request
self._page = getattr(response, self._name) or []
self._idx = 0
if not config:
request_config = {}
elif isinstance(config, dict):
request_config = copy.deepcopy(config)
else:
request_config = dict(config)
request_config['page_token'] = getattr(response, 'next_page_token')
self._config = request_config
self._page_size = request_config.get('page_size', len(self._page))
@property
def page(self) -> list[T]:
"""Returns the current page, which is a list of items.
The returned list of items is a subset of the entire list.
Usage:
.. code-block:: python
batch_jobs_pager = client.batches.list(config={'page_size': 5})
print(f"first page: {batch_jobs_pager.page}")
# first page: [BatchJob(name='projects/./locations/./batchPredictionJobs/1
"""
return self._page
@property
def name(self) -> str:
"""Returns the type of paged item (for example, ``batch_jobs``).
Usage:
.. code-block:: python
batch_jobs_pager = client.batches.list(config={'page_size': 5})
print(f"name: {batch_jobs_pager.name}")
# name: batch_jobs
"""
return self._name
@property
def page_size(self) -> int:
"""Returns the length of the page fetched each time by this pager.
The number of items in the page is less than or equal to the page length.
Usage:
.. code-block:: python
batch_jobs_pager = client.batches.list(config={'page_size': 5})
print(f"page_size: {batch_jobs_pager.page_size}")
# page_size: 5
"""
return self._page_size
@property
def config(self) -> dict[str, Any]:
"""Returns the configuration when making the API request for the next page.
A configuration is a set of optional parameters and arguments that can be
used to customize the API request. For example, the ``page_token`` parameter
contains the token to request the next page.
Usage:
.. code-block:: python
batch_jobs_pager = client.batches.list(config={'page_size': 5})
print(f"config: {batch_jobs_pager.config}")
# config: {'page_size': 5, 'page_token': 'AMEw9yO5jnsGnZJLHSKDFHJJu'}
"""
return self._config
def __len__(self) -> int:
"""Returns the total number of items in the current page."""
return len(self.page)
def __getitem__(self, index: int) -> T:
"""Returns the item at the given index."""
return self.page[index]
def _init_next_page(self, response: Any) -> None:
"""Initializes the next page from the response.
This is an internal method that should be called by subclasses after
fetching the next page.
Args:
response: The response object from the API request.
"""
self.__init__(self.name, self._request, response, self.config)
class Pager(_BasePager[T]):
"""Pager class for iterating through paginated results."""
def __next__(self) -> T:
"""Returns the next item."""
if self._idx >= len(self):
try:
self.next_page()
except IndexError:
raise StopIteration
item = self.page[self._idx]
self._idx += 1
return item
def __iter__(self) -> Iterator[T]:
"""Returns an iterator over the items."""
self._idx = 0
return self
def next_page(self) -> list[T]:
"""Fetches the next page of items. This makes a new API request.
Usage:
.. code-block:: python
batch_jobs_pager = client.batches.list(config={'page_size': 5})
print(f"current page: {batch_jobs_pager.page}")
batch_jobs_pager.next_page()
print(f"next page: {batch_jobs_pager.page}")
# current page: [BatchJob(name='projects/.../batchPredictionJobs/1
# next page: [BatchJob(name='projects/.../batchPredictionJobs/6
"""
if not self.config.get('page_token'):
raise IndexError('No more pages to fetch.')
response = self._request(config=self.config)
self._init_next_page(response)
return self.page
class AsyncPager(_BasePager[T]):
"""AsyncPager class for iterating through paginated results."""
def __init__(
self,
name: PagedItem,
request: Callable[Any, Awaitable[Any]],
response: Any,
config: Any,
):
super().__init__(name, request, response, config)
def __aiter__(self) -> AsyncIterator[T]:
"""Returns an async iterator over the items."""
self._idx = 0
return self
async def __anext__(self) -> Awaitable[T]:
"""Returns the next item asynchronously."""
if self._idx >= len(self):
try:
await self.next_page()
except IndexError:
raise StopAsyncIteration
item = self.page[self._idx]
self._idx += 1
return item
async def next_page(self) -> list[T]:
"""Fetches the next page of items asynchronously.
This makes a new API request.
Returns:
The next page of items.
Raises:
IndexError: No more pages to fetch.
Usage:
.. code-block:: python
batch_jobs_pager = await client.aio.batches.list(config={'page_size': 5})
print(f"current page: {batch_jobs_pager.page}")
await batch_jobs_pager.next_page()
print(f"next page: {batch_jobs_pager.page}")
# current page: [BatchJob(name='projects/.../batchPredictionJobs/1
# next page: [BatchJob(name='projects/.../batchPredictionJobs/6
"""
if not self.config.get('page_token'):
raise IndexError('No more pages to fetch.')
response = await self._request(config=self.config)
self._init_next_page(response)
return self.page
