aboutsummaryrefslogtreecommitdiff
path: root/.venv/lib/python3.12/site-packages/google/genai/pagers.py
diff options
context:
space:
mode:
Diffstat (limited to '.venv/lib/python3.12/site-packages/google/genai/pagers.py')
-rw-r--r--.venv/lib/python3.12/site-packages/google/genai/pagers.py245
1 files changed, 245 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/google/genai/pagers.py b/.venv/lib/python3.12/site-packages/google/genai/pagers.py
new file mode 100644
index 00000000..a90465c3
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/google/genai/pagers.py
@@ -0,0 +1,245 @@
+# 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