How to Implement Keyset Pagination in Python (Seek Method)

from dataclasses import dataclass
from typing import List, Optional


@dataclass
class Row:
    id: int
    name: str

    def __lt__(self, other: "Row") -> bool:
        return (self.id, self.name) < (other.id, other.name)


class MockKeysetPaginator:
    """Pagination using keyset (seek) method instead of OFFSET."""

    def __init__(self, rows: List[Row], page_size: int = 3):
        self.rows = sorted(rows)
        self.page_size = page_size

    def fetch_page(self, after: Optional[Row] = None) -> List[Row]:
        """Return page of rows that come after the given key."""
        if after is None:
            start_index = 0
        else:
            start_index = next(
                (i for i, row in enumerate(self.rows) if row > after),
                len(self.rows),
            )
        end_index = start_index + self.page_size
        return self.rows[start_index:end_index]


if __name__ == "__main__":
    data = [
        Row(3, "alice"),
        Row(1, "bob"),
        Row(2, "carol"),
        Row(5, "dave"),
        Row(4, "eve"),
        Row(6, "frank"),
        Row(7, "grace"),
    ]
    paginator = MockKeysetPaginator(data, page_size=3)

    cursor = None
    total_pages = 0
    while True:
        page = paginator.fetch_page(cursor)
        if not page:
            break
        total_pages += 1
        print(f"Page {total_pages}: {[row.name for row in page]}")
        cursor = page[-1]
    print(f"Total pages fetched: {total_pages}")