from __future__ import annotations from collections.abc import Callable from dataclasses import dataclass from math import ceil from typing import Any, Generic, TypeVar from sqlalchemy import func, select from sqlalchemy.orm import Session from sqlalchemy.sql import Select T = TypeVar("T") U = TypeVar("U") DEFAULT_PAGE = 1 DEFAULT_PAGE_SIZE = 20 MAX_PAGE_SIZE = 100 @dataclass(frozen=True) class PageParams: page: int page_size: int @property def offset(self) -> int: return (self.page - 1) * self.page_size @dataclass(frozen=True) class PageResult(Generic[T]): items: list[T] total: int page: int page_size: int @property def total_pages(self) -> int: if self.total <= 0: return 0 return ceil(self.total / self.page_size) @property def has_next(self) -> bool: return self.page < self.total_pages @property def has_previous(self) -> bool: return self.page > 1 and self.total_pages > 0 def map(self, mapper: Callable[[T], U]) -> PageResult[U]: return PageResult( items=[mapper(item) for item in self.items], total=self.total, page=self.page, page_size=self.page_size, ) def normalize_page_params( page: int | None, page_size: int | None, *, max_page_size: int = MAX_PAGE_SIZE, ) -> PageParams: normalized_page = max(DEFAULT_PAGE, int(page or DEFAULT_PAGE)) normalized_page_size = max(1, int(page_size or DEFAULT_PAGE_SIZE)) normalized_page_size = min(normalized_page_size, max_page_size) return PageParams(page=normalized_page, page_size=normalized_page_size) def paginate_select( db: Session, stmt: Select[Any], *, page: int | None, page_size: int | None, max_page_size: int = MAX_PAGE_SIZE, unique: bool = False, ) -> PageResult[Any]: params = normalize_page_params(page, page_size, max_page_size=max_page_size) count_stmt = select(func.count()).select_from(stmt.order_by(None).subquery()) total = int(db.scalar(count_stmt) or 0) page_stmt = stmt.limit(params.page_size).offset(params.offset) scalars = db.execute(page_stmt).scalars() if unique: scalars = scalars.unique() return PageResult( items=list(scalars.all()), total=total, page=params.page, page_size=params.page_size, )