"""Pagination support for KittyCAD API with OpenAI-style auto-iteration."""
from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, Type, TypeVar
from pydantic import BaseModel
T = TypeVar("T", bound=BaseModel)
def _validate_page(page: Optional[BaseModel], seen_tokens: set[str]) -> Optional[str]:
"""Validate a complete page before yielding any of its items."""
if not isinstance(page, BaseModel) or "items" not in page.model_fields_set:
raise ValueError("Paginated response must include items")
if not isinstance(getattr(page, "items", None), list):
raise ValueError("Paginated response items must be a list")
# The API schema permits an omitted or null cursor on the final page.
next_page = getattr(page, "next_page", None)
if next_page is None:
return None
if not isinstance(next_page, str) or not next_page.strip():
raise ValueError("Paginated response has an invalid next_page token")
if next_page in seen_tokens:
raise ValueError("Paginated response repeated a page token")
seen_tokens.add(next_page)
return next_page
[docs]
class SyncPageIterator:
"""Synchronous iterator for paginated API responses.
Provides OpenAI-style auto-pagination that handles page tokens automatically.
"""
def __init__(
self,
page_fetcher: Callable[..., BaseModel],
initial_kwargs: Dict[str, Any],
item_type: Optional[Type[T]] = None,
):
"""Initialize the sync page iterator.
Args:
page_fetcher: Function to fetch a page (e.g., client.api_calls.list_api_calls)
initial_kwargs: Initial arguments for the first request
item_type: Type of items being paginated
"""
self._page_fetcher = page_fetcher
self._initial_kwargs = initial_kwargs
self._item_type = item_type
self._current_page_token: Optional[str] = None
self._exhausted = False
def __iter__(self) -> Iterator[T]:
"""Return iterator that yields individual items across all pages."""
# Reset state for new iteration
self._current_page_token = None
self._exhausted = False
kwargs = self._initial_kwargs.copy()
seen_tokens: set[str] = set()
if kwargs.get("page_token") is not None:
seen_tokens.add(kwargs["page_token"])
while not self._exhausted:
# Add page token if we have one
if self._current_page_token:
kwargs["page_token"] = self._current_page_token
# Note: Don't remove page_token from kwargs if it exists in initial_kwargs
# This allows users to explicitly start pagination from a specific token
# Fetch the page
page = self._page_fetcher(**kwargs)
# Fail before yielding this page if its continuation is malformed.
next_page_token = _validate_page(page, seen_tokens)
for item in getattr(page, "items"):
yield item
# Empty pages can still have a continuation.
if next_page_token is not None:
self._current_page_token = next_page_token
else:
self._exhausted = True
break
[docs]
class AsyncPageIterator:
"""Asynchronous iterator for paginated API responses.
Provides OpenAI-style async auto-pagination that handles page tokens automatically.
"""
def __init__(
self,
page_fetcher: Callable[..., Any], # Returns Awaitable[BaseModel]
initial_kwargs: Dict[str, Any],
item_type: Optional[Type[T]] = None,
):
"""Initialize the async page iterator.
Args:
page_fetcher: Async function to fetch a page (e.g., client.api_calls.list_api_calls)
initial_kwargs: Initial arguments for the first request
item_type: Type of items being paginated
"""
self._page_fetcher = page_fetcher
self._initial_kwargs = initial_kwargs
self._item_type = item_type
self._current_page_token: Optional[str] = None
self._exhausted = False
def __aiter__(self) -> AsyncIterator[T]:
"""Return async iterator that yields individual items across all pages."""
# Reset state for new iteration
self._current_page_token = None
self._exhausted = False
return self._async_iter()
async def _async_iter(self) -> AsyncIterator[T]:
"""Internal async iterator implementation."""
kwargs = self._initial_kwargs.copy()
seen_tokens: set[str] = set()
if kwargs.get("page_token") is not None:
seen_tokens.add(kwargs["page_token"])
while not self._exhausted:
# Add page token if we have one
if self._current_page_token:
kwargs["page_token"] = self._current_page_token
# Note: Don't remove page_token from kwargs if it exists in initial_kwargs
# This allows users to explicitly start pagination from a specific token
# Fetch the page
page = await self._page_fetcher(**kwargs)
# Fail before yielding this page if its continuation is malformed.
next_page_token = _validate_page(page, seen_tokens)
for item in getattr(page, "items"):
yield item
# Empty pages can still have a continuation.
if next_page_token is not None:
self._current_page_token = next_page_token
else:
self._exhausted = True
break
[docs]
def create_sync_page_iterator(
page_fetcher: Callable[..., BaseModel],
kwargs: Dict[str, Any],
item_type: Optional[Type[T]] = None,
) -> SyncPageIterator:
"""Create a synchronous page iterator.
Args:
page_fetcher: Function that fetches a page
kwargs: Arguments to pass to the page fetcher
item_type: Type of individual items
Returns:
SyncPageIterator that can be iterated to get all items
"""
return SyncPageIterator(page_fetcher, kwargs, item_type)
[docs]
def create_async_page_iterator(
page_fetcher: Callable[..., Any], # Returns Awaitable[BaseModel]
kwargs: Dict[str, Any],
item_type: Optional[Type[T]] = None,
) -> AsyncPageIterator:
"""Create an asynchronous page iterator.
Args:
page_fetcher: Async function that fetches a page
kwargs: Arguments to pass to the page fetcher
item_type: Type of individual items
Returns:
AsyncPageIterator that can be async iterated to get all items
"""
return AsyncPageIterator(page_fetcher, kwargs, item_type)