diff options
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.py | 245 |
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 |