paginator.py

# AUTO-GENERATED FILE - DO NOT EDIT
# This file was automatically generated by the XDK build tool.
# Any manual changes will be overwritten on the next generation.
"""
Cursor-based pagination utilities for the X API SDK.

This module provides a Cursor class for elegant pagination support across
all API clients. The Cursor enables easy iteration over paginated results
using both .pages() and .items() methods with proper type safety.

Generated automatically - do not edit manually.
"""

from typing import (
    Iterator,
    Any,
    Optional,
    Callable,
    TypeVar,
    Generic,
    overload,
    Protocol,
    runtime_checkable,
)
import inspect
import types

# Type variable for response type
ResponseType = TypeVar("ResponseType")


@runtime_checkable
class PaginatableMethod(Protocol[ResponseType]):
    """
    Protocol for methods that support pagination.

    A method is considered paginatable if it accepts pagination_token
    and/or max_results parameters.
    """


    def __call__(
        self,
        *args,
        pagination_token: Optional[str] = None,
        max_results: Optional[int] = None,
        **kwargs,
    ) -> ResponseType: ...


class PaginationError(TypeError):
    """Raised when a non-paginatable method is passed to Cursor."""

    pass


def _validate_paginatable_method(method: Callable) -> None:
    """
    Validate that a method supports pagination parameters.
    Raises PaginationError if the method doesn't support pagination.
    """
    sig = inspect.signature(method)
    params = sig.parameters
    # Check if method has pagination_token, next_token, or max_results parameters
    has_pagination_token = "pagination_token" in params
    has_next_token = "next_token" in params
    has_max_results = "max_results" in params
    if not (has_pagination_token or has_next_token or has_max_results):
        raise PaginationError(
            f"Method '{method.__name__}' does not support pagination. "
            f"Paginatable methods must accept 'pagination_token', 'next_token', and/or 'max_results' parameters. "
            f"Available parameters: {list(params.keys())}"
        )


class Cursor(Generic[ResponseType]):
    """
    A cursor for handling pagination in API calls.

    This class wraps any API method and provides elegant pagination
    through .items() and .pages() methods. The cursor preserves the
    return type of the wrapped method for proper type hints.

    Example usage:
        .. code-block:: python

            # For methods requiring arguments
            cursor = Cursor(client.users.get_users_blocking, "user_id", max_results=100)
            for user in cursor.items(250):  # user is properly typed
                print(user.name)

            # For methods with query parameters
            cursor = Cursor(client.posts.search_posts_recent, "python", max_results=50)
            for page in cursor.pages(5):  # page is SearchResponse type
                print(f"Got {len(page.data)} posts")
    """


    def __init__(
        self, method: PaginatableMethod[ResponseType], *args: Any, **kwargs: Any
    ):
        """
        Initialize the cursor.
        Args:
            method: The API method to call for each page (must support pagination)
            *args: Positional arguments to pass to the API method
            **kwargs: Keyword arguments to pass to the API method (excluding pagination params)
        Raises:
            PaginationError: If the method doesn't support pagination parameters
        """
        # Validate method supports pagination at runtime
        _validate_paginatable_method(method)
        self.method = method
        self.args = args
        self.params = kwargs.copy()


    def items(self, limit: Optional[int] = None) -> Iterator[Any]:
        """
        Iterate over individual items from paginated responses.
        Args:
            limit: Maximum number of items to return (None for unlimited)
        Yields:
            Individual items from the API responses
        """
        return _ItemIterator(self, limit)


    def pages(self, limit: Optional[int] = None) -> Iterator[ResponseType]:
        """
        Iterate over pages of responses.
        Args:
            limit: Maximum number of pages to return (None for unlimited)
        Yields:
            Page responses from the API (same type as the wrapped method returns)
        """
        return _PageIterator(self, limit)


# Factory function for better type inference


@overload


def cursor(
    method: PaginatableMethod[ResponseType], *args: Any, **kwargs: Any
) -> Cursor[ResponseType]: ...


def cursor(
    method: PaginatableMethod[ResponseType], *args: Any, **kwargs: Any
) -> Cursor[ResponseType]:
    """
    Create a cursor with proper type inference and validation.
    This factory function helps with type inference so you get proper
    type hints for the response type, and validates that the method
    supports pagination at both static analysis and runtime.
    Args:
        method: The API method to wrap (must support pagination)
        *args: Positional arguments to pass to the method
        **kwargs: Keyword arguments to pass to the method
    Returns:
        A properly typed Cursor instance
    Raises:
        PaginationError: If the method doesn't support pagination parameters
    Example:
        .. code-block:: python
            # Type is inferred as Cursor[GetUsersResponse]
            users_cursor = cursor(client.users.get_users_blocking, "user_id", max_results=100)
            # page is typed as GetUsersResponse
            for page in users_cursor.pages(5):
                print(len(page.data))
            # For search methods
            search_cursor = cursor(client.posts.search_posts_recent, "python", max_results=50)
            for tweet in search_cursor.items(100):
                print(tweet.text)
    """
    return Cursor(method, *args, **kwargs)


class _PageIterator(Generic[ResponseType], Iterator[ResponseType]):
    """Internal iterator for pages."""


    def __init__(self, cursor: Cursor[ResponseType], limit: Optional[int]):
        self.cursor = cursor
        self.limit = limit
        self.count = 0
        self.exhausted = False
        self._next_token: Optional[str] = None
        self._generator: Optional[Iterator[ResponseType]] = None
        self._is_generator_method = False


    def __iter__(self) -> Iterator[ResponseType]:
        return self


    def __next__(self) -> ResponseType:
        if self.exhausted or (self.limit is not None and self.count >= self.limit):
            raise StopIteration
        # Check if method returns a generator (first call only)
        if self._generator is None:
            # Prepare params for this request
            params = self.cursor.params.copy()
            # Don't pass pagination_token/next_token on first call for generator methods
            # as they handle pagination internally
            # Make the API call with both positional and keyword arguments
            result = self.cursor.method(*self.cursor.args, **params)
            # Check if the method returned a generator/iterator
            # Paginated methods now return generators that yield pages
            # Primary check: isinstance with GeneratorType
            is_generator = isinstance(result, types.GeneratorType)
            # Fallback: check if it's a generator-like object (has __iter__ and __next__)
            # but exclude common iterables like str, list, dict, tuple
            if not is_generator:
                if (
                    hasattr(result, "__iter__")
                    and hasattr(result, "__next__")
                    and not isinstance(result, (str, bytes, list, dict, tuple, set))
                ):
                    try:
                        # If iter() returns itself, it's an iterator
                        iter_result = iter(result)
                        if iter_result is result:
                            is_generator = True
                    except (TypeError, AttributeError):
                        pass
            if is_generator:
                self._is_generator_method = True
                self._generator = iter(result)
                # For generator methods, we just yield from the generator
                try:
                    response = next(self._generator)
                    self.count += 1
                    return response
                except StopIteration:
                    self.exhausted = True
                    raise StopIteration
            else:
                # Old behavior: method returns a single response
                self._is_generator_method = False
                # Extract next token AFTER returning this response
                extracted_token = self._extract_next_token(result)
                self._next_token = extracted_token
                # Check if we're done - if no token was extracted, we've reached the end
                if not extracted_token:
                    self.exhausted = True
                self.count += 1
                return result
        # If we have a generator, continue yielding from it
        if self._is_generator_method:
            try:
                response = next(self._generator)
                self.count += 1
                return response
            except StopIteration:
                self.exhausted = True
                raise StopIteration
        # Old behavior: method returns single responses, need to call again
        # Prepare params for this request
        params = self.cursor.params.copy()
        if self._next_token:
            # Try different pagination parameter names
            if self._supports_pagination_token():
                params["pagination_token"] = self._next_token
            elif self._supports_next_token():
                params["next_token"] = self._next_token
        # Make the API call with both positional and keyword arguments
        response = self.cursor.method(*self.cursor.args, **params)
        # Extract next token AFTER returning this response
        # This allows us to check if there are more pages
        extracted_token = self._extract_next_token(response)
        self._next_token = extracted_token
        # Check if we're done - if no token was extracted, we've reached the end
        if not extracted_token:
            self.exhausted = True
        self.count += 1
        return response


    def _supports_pagination_token(self) -> bool:
        """Check if the method supports pagination_token parameter."""
        sig = inspect.signature(self.cursor.method)
        return "pagination_token" in sig.parameters


    def _supports_next_token(self) -> bool:
        """Check if the method supports next_token parameter."""
        sig = inspect.signature(self.cursor.method)
        return "next_token" in sig.parameters


    def _extract_next_token(self, response: ResponseType) -> Optional[str]:
        """Extract the next_token from the response."""
        # First, try to convert Pydantic model to dict if possible
        response_dict = None
        if hasattr(response, "model_dump"):
            try:
                # model_dump includes both defined fields and extra fields
                response_dict = response.model_dump()
            except (AttributeError, TypeError):
                pass
        # Also try to get extra fields directly from Pydantic
        if hasattr(response, "__pydantic_extra__"):
            try:
                extra = response.__pydantic_extra__
                if extra and isinstance(extra, dict):
                    meta = extra.get("meta")
                    if meta and isinstance(meta, dict):
                        token = meta.get("next_token")
                        if token:
                            return str(token)
                    # Also check for next_token at root level in extra
                    token = extra.get("next_token")
                    if token:
                        return str(token)
            except (AttributeError, TypeError, KeyError):
                pass
        # Pattern 1: response.meta.next_token (most common)
        try:
            # Try Pydantic model attributes first
            if hasattr(response, "meta") and response.meta is not None:
                meta = response.meta
                # If meta is a Pydantic model, try to dump it
                if hasattr(meta, "model_dump"):
                    try:
                        meta_dict = meta.model_dump()
                        token = meta_dict.get("next_token")
                        if token:
                            return str(token)
                    except (AttributeError, TypeError):
                        pass
                # Otherwise try attribute access
                if hasattr(meta, "next_token"):
                    token = getattr(meta, "next_token", None)
                    if token:
                        return str(token)
                # If meta is a dict, access it directly
                if isinstance(meta, dict):
                    token = meta.get("next_token")
                    if token:
                        return str(token)
        except (AttributeError, TypeError):
            pass
        # Try dict access if we have a dict
        if response_dict:
            try:
                meta = response_dict.get("meta")
                if meta and isinstance(meta, dict):
                    token = meta.get("next_token")
                    if token:
                        return str(token)
            except (AttributeError, TypeError, KeyError):
                pass
        # Pattern 2: response.next_token (some APIs)
        try:
            if hasattr(response, "next_token"):
                token = getattr(response, "next_token", None)
                if token:
                    return str(token)
        except (AttributeError, TypeError):
            pass
        # Try dict access for next_token at root level
        if response_dict:
            try:
                token = response_dict.get("next_token")
                if token:
                    return str(token)
            except (AttributeError, TypeError, KeyError):
                pass
        return None


class _ItemIterator(Iterator[Any]):
    """Internal iterator for individual items."""


    def __init__(self, cursor: Cursor, limit: Optional[int]):
        self.page_iterator = _PageIterator(cursor, None)  # No page limit for items
        self.limit = limit
        self.count = 0
        self._current_items = []
        self._current_index = 0


    def __iter__(self) -> Iterator[Any]:
        return self


    def __next__(self) -> Any:
        if self.limit is not None and self.count >= self.limit:
            raise StopIteration
        # If we've exhausted current items, get next page
        while self._current_index >= len(self._current_items):
            try:
                page = next(self.page_iterator)
                # Try common patterns for data arrays
                items = self._extract_items(page)
                self._current_items = items or []
                self._current_index = 0
                # If the page has no items, continue to next page
                if not self._current_items:
                    continue
            except StopIteration:
                raise StopIteration
        # Return current item and advance
        item = self._current_items[self._current_index]
        self._current_index += 1
        self.count += 1
        return item


    def _extract_items(self, response: Any) -> Optional[list]:
        """Extract items from response, trying common field names."""
        # First, try to convert Pydantic model to dict if possible
        response_dict = None
        if hasattr(response, "model_dump"):
            try:
                # model_dump includes both defined fields and extra fields
                response_dict = response.model_dump()
            except (AttributeError, TypeError):
                pass
        # Also try to get extra fields directly from Pydantic
        if hasattr(response, "__pydantic_extra__"):
            try:
                extra = response.__pydantic_extra__
                if extra and isinstance(extra, dict):
                    for field_name in ["data", "results", "items"]:
                        items = extra.get(field_name)
                        if isinstance(items, list):
                            return items
            except (AttributeError, TypeError, KeyError):
                pass
        # Try common field names for data arrays
        for field_name in ["data", "results", "items"]:
            # Try attribute access first (for Pydantic models)
            try:
                if hasattr(response, field_name):
                    items = getattr(response, field_name, None)
                    if isinstance(items, list):
                        return items
            except (AttributeError, TypeError):
                pass
            # Try dict access if we have a dict
            if response_dict:
                try:
                    items = response_dict.get(field_name)
                    if isinstance(items, list):
                        return items
                except (AttributeError, TypeError, KeyError):
                    continue
        return None