Skip to content

Commit 985ec36

Browse files
committed
feat: add PaginationHelper utility class
1 parent e878120 commit 985ec36

File tree

3 files changed

+234
-0
lines changed

3 files changed

+234
-0
lines changed

src/gradient/_utils/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
RateLimiter as RateLimiter,
3737
BatchProcessor as BatchProcessor,
3838
DataExporter as DataExporter,
39+
PaginationHelper as PaginationHelper,
3940
)
4041
from ._compat import (
4142
get_args as get_args,

src/gradient/_utils/_utils.py

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -715,6 +715,129 @@ def export_csv(self, data: Any, file_path: str | None = None, headers: list[str]
715715
output.close()
716716

717717

718+
# Pagination Classes
719+
class PaginationHelper:
720+
"""Helper for handling paginated API responses."""
721+
722+
def __init__(self, page_size: int = 20, max_pages: int | None = None) -> None:
723+
"""Initialize pagination helper.
724+
725+
Args:
726+
page_size: Number of items per page
727+
max_pages: Maximum number of pages to fetch (None for unlimited)
728+
"""
729+
self.page_size: int = page_size
730+
self.max_pages: int | None = max_pages
731+
732+
def paginate(self, fetch_func: Callable[[dict[str, Any]], Any], **kwargs: Any) -> list[Any]:
733+
"""Paginate through all results using the provided fetch function.
734+
735+
Args:
736+
fetch_func: Function that takes pagination params and returns response
737+
**kwargs: Additional parameters to pass to fetch_func
738+
739+
Returns:
740+
List of all items across all pages
741+
"""
742+
all_items = []
743+
page = 1
744+
745+
while self.max_pages is None or page <= self.max_pages:
746+
# Add pagination parameters
747+
params = kwargs.copy()
748+
params.update({
749+
"page": page,
750+
"per_page": self.page_size
751+
})
752+
753+
try:
754+
response = fetch_func(params)
755+
items = self._extract_items(response)
756+
757+
if not items:
758+
break # No more items
759+
760+
all_items.extend(items)
761+
762+
# Check if we got fewer items than requested (last page)
763+
if len(items) < self.page_size:
764+
break
765+
766+
page += 1
767+
768+
except Exception as e:
769+
# If it's a pagination error or no more pages, stop
770+
if self._is_pagination_end_error(e):
771+
break
772+
raise
773+
774+
return all_items
775+
776+
def _extract_items(self, response: Any) -> list[Any]:
777+
"""Extract items from API response."""
778+
# Handle different response formats
779+
if hasattr(response, 'data') and isinstance(response.data, list):
780+
return response.data
781+
elif hasattr(response, 'items') and isinstance(response.items, list):
782+
return response.items
783+
elif hasattr(response, 'results') and isinstance(response.results, list):
784+
return response.results
785+
elif isinstance(response, list):
786+
return response
787+
elif isinstance(response, dict):
788+
# Try common keys
789+
for key in ['data', 'items', 'results', 'objects']:
790+
if key in response and isinstance(response[key], list):
791+
return response[key]
792+
return []
793+
794+
def _is_pagination_end_error(self, error: Exception) -> bool:
795+
"""Check if error indicates end of pagination."""
796+
error_str = str(error).lower()
797+
return any(phrase in error_str for phrase in [
798+
'page not found',
799+
'invalid page',
800+
'no more pages',
801+
'pagination end'
802+
])
803+
804+
async def paginate_async(self, fetch_func: Callable[[dict[str, Any]], Any], **kwargs: Any) -> list[Any]:
805+
"""Async version of paginate."""
806+
import asyncio
807+
808+
all_items = []
809+
page = 1
810+
811+
while self.max_pages is None or page <= self.max_pages:
812+
# Add pagination parameters
813+
params = kwargs.copy()
814+
params.update({
815+
"page": page,
816+
"per_page": self.page_size
817+
})
818+
819+
try:
820+
response = await fetch_func(params)
821+
items = self._extract_items(response)
822+
823+
if not items:
824+
break
825+
826+
all_items.extend(items)
827+
828+
if len(items) < self.page_size:
829+
break
830+
831+
page += 1
832+
833+
except Exception as e:
834+
if self._is_pagination_end_error(e):
835+
break
836+
raise
837+
838+
return all_items
839+
840+
718841
# API Key Validation Functions
719842
def validate_api_key(api_key: str | None) -> bool:
720843
"""Validate an API key format.

tests/test_pagination_helper.py

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
"""Tests for pagination helper functionality."""
2+
3+
import pytest
4+
from gradient._utils import PaginationHelper
5+
6+
7+
class TestPaginationHelper:
8+
"""Test pagination helper functionality."""
9+
10+
def test_pagination_helper_basic(self):
11+
"""Test basic pagination functionality."""
12+
helper = PaginationHelper(page_size=2, max_pages=3)
13+
14+
# Mock responses with different pages
15+
responses = [
16+
{"data": ["item1", "item2"]}, # Page 1
17+
{"data": ["item3", "item4"]}, # Page 2
18+
{"data": ["item5"]}, # Page 3 (partial)
19+
]
20+
response_index = 0
21+
22+
def mock_fetch(params):
23+
nonlocal response_index
24+
if response_index < len(responses):
25+
result = responses[response_index]
26+
response_index += 1
27+
return result
28+
return {"data": []}
29+
30+
result = helper.paginate(mock_fetch)
31+
assert result == ["item1", "item2", "item3", "item4", "item5"]
32+
33+
def test_pagination_helper_max_pages(self):
34+
"""Test pagination with max pages limit."""
35+
helper = PaginationHelper(page_size=2, max_pages=2)
36+
37+
responses = [
38+
{"data": ["item1", "item2"]},
39+
{"data": ["item3", "item4"]},
40+
{"data": ["item5", "item6"]}, # Should not be reached
41+
]
42+
response_index = 0
43+
44+
def mock_fetch(params):
45+
nonlocal response_index
46+
if response_index < len(responses):
47+
result = responses[response_index]
48+
response_index += 1
49+
return result
50+
return {"data": []}
51+
52+
result = helper.paginate(mock_fetch)
53+
assert result == ["item1", "item2", "item3", "item4"]
54+
55+
def test_pagination_helper_empty_response(self):
56+
"""Test pagination with empty response."""
57+
helper = PaginationHelper(page_size=2)
58+
59+
def mock_fetch(params):
60+
return {"data": []}
61+
62+
result = helper.paginate(mock_fetch)
63+
assert result == []
64+
65+
def test_pagination_helper_different_response_formats(self):
66+
"""Test pagination with different response formats."""
67+
helper = PaginationHelper(page_size=2)
68+
69+
# Test different response formats
70+
test_cases = [
71+
{"data": ["item1", "item2"]},
72+
{"items": ["item3", "item4"]},
73+
{"results": ["item5", "item6"]},
74+
["item7", "item8"], # Direct list
75+
]
76+
77+
for i, expected_format in enumerate(test_cases):
78+
response_index = 0
79+
80+
def mock_fetch(params):
81+
nonlocal response_index
82+
if response_index == 0:
83+
response_index += 1
84+
return expected_format
85+
return {"data": []}
86+
87+
result = helper.paginate(mock_fetch)
88+
expected_items = expected_format if isinstance(expected_format, list) else expected_format.get(
89+
list(expected_format.keys())[0], []
90+
)
91+
assert result == expected_items
92+
93+
def test_pagination_helper_extract_items(self):
94+
"""Test item extraction from different response formats."""
95+
helper = PaginationHelper()
96+
97+
# Test various response formats
98+
test_responses = [
99+
type('MockResponse', (), {'data': ['item1']})(),
100+
type('MockResponse', (), {'items': ['item2']})(),
101+
type('MockResponse', (), {'results': ['item3']})(),
102+
['item4'],
103+
{'data': ['item5']},
104+
{'unknown': ['should_be_empty']},
105+
]
106+
107+
for i, response in enumerate(test_responses):
108+
items = helper._extract_items(response)
109+
expected_items = [f'item{i+1}'] if i < 5 else []
110+
assert items == expected_items

0 commit comments

Comments
 (0)