diff --git a/src/codex/pagination.py b/src/codex/pagination.py index 4798baa..b38b5f4 100644 --- a/src/codex/pagination.py +++ b/src/codex/pagination.py @@ -1,16 +1,27 @@ # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. -from typing import List, Generic, TypeVar, Optional +from typing import Any, List, Type, Generic, Mapping, TypeVar, Optional, cast from typing_extensions import override +from httpx import Response + +from ._utils import is_mapping +from ._models import BaseModel from ._base_client import BasePage, PageInfo, BaseSyncPage, BaseAsyncPage -__all__ = ["SyncMyOffsetPage", "AsyncMyOffsetPage"] +__all__ = [ + "SyncMyOffsetPageTopLevelArray", + "AsyncMyOffsetPageTopLevelArray", + "SyncOffsetPageEntries", + "AsyncOffsetPageEntries", +] + +_BaseModelT = TypeVar("_BaseModelT", bound=BaseModel) _T = TypeVar("_T") -class SyncMyOffsetPage(BaseSyncPage[_T], BasePage[_T], Generic[_T]): +class SyncMyOffsetPageTopLevelArray(BaseSyncPage[_T], BasePage[_T], Generic[_T]): items: List[_T] @override @@ -31,8 +42,17 @@ def next_page_info(self) -> Optional[PageInfo]: return PageInfo(params={"offset": current_count}) + @classmethod + def build(cls: Type[_BaseModelT], *, response: Response, data: object) -> _BaseModelT: # noqa: ARG003 + return cls.construct( + None, + **{ + **(cast(Mapping[str, Any], data) if is_mapping(data) else {"items": data}), + }, + ) + -class AsyncMyOffsetPage(BaseAsyncPage[_T], BasePage[_T], Generic[_T]): +class AsyncMyOffsetPageTopLevelArray(BaseAsyncPage[_T], BasePage[_T], Generic[_T]): items: List[_T] @override @@ -52,3 +72,90 @@ def next_page_info(self) -> Optional[PageInfo]: current_count = offset + length return PageInfo(params={"offset": current_count}) + + @classmethod + def build(cls: Type[_BaseModelT], *, response: Response, data: object) -> _BaseModelT: # noqa: ARG003 + return cls.construct( + None, + **{ + **(cast(Mapping[str, Any], data) if is_mapping(data) else {"items": data}), + }, + ) + + +class SyncOffsetPageEntries(BaseSyncPage[_T], BasePage[_T], Generic[_T]): + entries: List[_T] + total_count: Optional[int] = None + + @override + def _get_page_items(self) -> List[_T]: + entries = self.entries + if not entries: + return [] + return entries + + @override + def next_page_info(self) -> Optional[PageInfo]: + offset = self._options.params.get("offset") or 0 + if not isinstance(offset, int): + raise ValueError(f'Expected "offset" param to be an integer but got {offset}') + + length = len(self._get_page_items()) + current_count = offset + length + + total_count = self.total_count + if total_count is None: + return None + + if current_count < total_count: + return PageInfo(params={"offset": current_count}) + + return None + + @classmethod + def build(cls: Type[_BaseModelT], *, response: Response, data: object) -> _BaseModelT: # noqa: ARG003 + return cls.construct( + None, + **{ + **(cast(Mapping[str, Any], data) if is_mapping(data) else {"entries": data}), + }, + ) + + +class AsyncOffsetPageEntries(BaseAsyncPage[_T], BasePage[_T], Generic[_T]): + entries: List[_T] + total_count: Optional[int] = None + + @override + def _get_page_items(self) -> List[_T]: + entries = self.entries + if not entries: + return [] + return entries + + @override + def next_page_info(self) -> Optional[PageInfo]: + offset = self._options.params.get("offset") or 0 + if not isinstance(offset, int): + raise ValueError(f'Expected "offset" param to be an integer but got {offset}') + + length = len(self._get_page_items()) + current_count = offset + length + + total_count = self.total_count + if total_count is None: + return None + + if current_count < total_count: + return PageInfo(params={"offset": current_count}) + + return None + + @classmethod + def build(cls: Type[_BaseModelT], *, response: Response, data: object) -> _BaseModelT: # noqa: ARG003 + return cls.construct( + None, + **{ + **(cast(Mapping[str, Any], data) if is_mapping(data) else {"entries": data}), + }, + )