Skip to content

Commit e733095

Browse files
committed
feat: add PaginationHelper utility class
1 parent 76392b5 commit e733095

File tree

3 files changed

+351
-0
lines changed

3 files changed

+351
-0
lines changed

src/gradient/_utils/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
get_required_header as get_required_header,
3030
maybe_coerce_boolean as maybe_coerce_boolean,
3131
maybe_coerce_integer as maybe_coerce_integer,
32+
PaginationHelper as PaginationHelper,
3233
)
3334
from ._compat import (
3435
get_args as get_args,

src/gradient/_utils/_utils.py

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -419,3 +419,126 @@ def json_safe(data: object) -> object:
419419
return data.isoformat()
420420

421421
return data
422+
423+
424+
# Pagination Classes
425+
class PaginationHelper:
426+
"""Helper for handling paginated API responses."""
427+
428+
def __init__(self, page_size: int = 20, max_pages: int | None = None) -> None:
429+
"""Initialize pagination helper.
430+
431+
Args:
432+
page_size: Number of items per page
433+
max_pages: Maximum number of pages to fetch (None for unlimited)
434+
"""
435+
self.page_size: int = page_size
436+
self.max_pages: int | None = max_pages
437+
438+
def paginate(self, fetch_func: Callable[[dict[str, Any]], Any], **kwargs: Any) -> list[Any]:
439+
"""Paginate through all results using the provided fetch function.
440+
441+
Args:
442+
fetch_func: Function that takes pagination params and returns response
443+
**kwargs: Additional parameters to pass to fetch_func
444+
445+
Returns:
446+
List of all items across all pages
447+
"""
448+
all_items = []
449+
page = 1
450+
451+
while self.max_pages is None or page <= self.max_pages:
452+
# Add pagination parameters
453+
params = kwargs.copy()
454+
params.update({
455+
"page": page,
456+
"per_page": self.page_size
457+
})
458+
459+
try:
460+
response = fetch_func(params)
461+
items = self._extract_items(response)
462+
463+
if not items:
464+
break # No more items
465+
466+
all_items.extend(items)
467+
468+
# Check if we got fewer items than requested (last page)
469+
if len(items) < self.page_size:
470+
break
471+
472+
page += 1
473+
474+
except Exception as e:
475+
# If it's a pagination error or no more pages, stop
476+
if self._is_pagination_end_error(e):
477+
break
478+
raise
479+
480+
return all_items
481+
482+
def _extract_items(self, response: Any) -> list[Any]:
483+
"""Extract items from API response."""
484+
# Handle different response formats
485+
if hasattr(response, 'data') and isinstance(response.data, list):
486+
return response.data
487+
elif hasattr(response, 'items') and isinstance(response.items, list):
488+
return response.items
489+
elif hasattr(response, 'results') and isinstance(response.results, list):
490+
return response.results
491+
elif isinstance(response, list):
492+
return response
493+
elif isinstance(response, dict):
494+
# Try common keys
495+
for key in ['data', 'items', 'results', 'objects']:
496+
if key in response and isinstance(response[key], list):
497+
return response[key]
498+
return []
499+
500+
def _is_pagination_end_error(self, error: Exception) -> bool:
501+
"""Check if error indicates end of pagination."""
502+
error_str = str(error).lower()
503+
return any(phrase in error_str for phrase in [
504+
'page not found',
505+
'invalid page',
506+
'no more pages',
507+
'pagination end'
508+
])
509+
510+
async def paginate_async(self, fetch_func: Callable[[dict[str, Any]], Any], **kwargs: Any) -> list[Any]:
511+
"""Async version of paginate."""
512+
import asyncio
513+
514+
all_items = []
515+
page = 1
516+
517+
while self.max_pages is None or page <= self.max_pages:
518+
# Add pagination parameters
519+
params = kwargs.copy()
520+
params.update({
521+
"page": page,
522+
"per_page": self.page_size
523+
})
524+
525+
try:
526+
response = await fetch_func(params)
527+
items = self._extract_items(response)
528+
529+
if not items:
530+
break
531+
532+
all_items.extend(items)
533+
534+
if len(items) < self.page_size:
535+
break
536+
537+
page += 1
538+
539+
except Exception as e:
540+
if self._is_pagination_end_error(e):
541+
break
542+
raise
543+
544+
return all_items

tests/test_pagination_helper.py

Lines changed: 227 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,227 @@
1+
"""Tests for PaginationHelper utility class."""
2+
3+
from __future__ import annotations
4+
5+
import pytest
6+
from unittest.mock import Mock, call
7+
8+
from gradient._utils import PaginationHelper
9+
10+
11+
class TestPaginationHelper:
12+
"""Test cases for PaginationHelper class."""
13+
14+
def test_init(self) -> None:
15+
"""Test PaginationHelper initialization."""
16+
helper = PaginationHelper(page_size=50, max_pages=10)
17+
assert helper.page_size == 50
18+
assert helper.max_pages == 10
19+
20+
# Test defaults
21+
helper_default = PaginationHelper()
22+
assert helper_default.page_size == 20
23+
assert helper_default.max_pages is None
24+
25+
def test_paginate_basic(self) -> None:
26+
"""Test basic pagination functionality."""
27+
helper = PaginationHelper(page_size=2, max_pages=3)
28+
29+
# Mock fetch function that returns different data for each page
30+
call_count = 0
31+
def mock_fetch(params):
32+
nonlocal call_count
33+
call_count += 1
34+
page = params.get('page', 1)
35+
per_page = params.get('per_page', 2)
36+
37+
if call_count == 1:
38+
return {'data': ['item1', 'item2']}
39+
elif call_count == 2:
40+
return {'data': ['item3', 'item4']}
41+
elif call_count == 3:
42+
return {'data': ['item5']} # Last page with fewer items
43+
else:
44+
return {'data': []} # No more data
45+
46+
result = helper.paginate(mock_fetch)
47+
48+
assert result == ['item1', 'item2', 'item3', 'item4', 'item5']
49+
assert call_count == 3
50+
51+
def test_paginate_with_kwargs(self) -> None:
52+
"""Test pagination with additional kwargs."""
53+
helper = PaginationHelper(page_size=1)
54+
55+
def mock_fetch(params):
56+
# Should include both pagination params and custom params
57+
assert 'page' in params
58+
assert 'per_page' in params
59+
assert params['custom_param'] == 'value'
60+
return {'data': [f'item{params["page"]}']}
61+
62+
# Mock to return only one page
63+
call_count = 0
64+
def mock_fetch_single(params):
65+
nonlocal call_count
66+
call_count += 1
67+
if call_count == 1:
68+
return {'data': ['item1']}
69+
return {'data': []}
70+
71+
result = helper.paginate(mock_fetch_single, custom_param='value')
72+
assert result == ['item1']
73+
74+
def test_paginate_max_pages_limit(self) -> None:
75+
"""Test that max_pages limits the number of pages fetched."""
76+
helper = PaginationHelper(page_size=1, max_pages=2)
77+
78+
call_count = 0
79+
def mock_fetch(params):
80+
nonlocal call_count
81+
call_count += 1
82+
return {'data': [f'item{call_count}']}
83+
84+
result = helper.paginate(mock_fetch)
85+
assert result == ['item1', 'item2']
86+
assert call_count == 2
87+
88+
def test_paginate_different_response_formats(self) -> None:
89+
"""Test pagination with different response formats."""
90+
helper = PaginationHelper(page_size=1)
91+
92+
# Test different response formats
93+
responses = [
94+
{'data': ['item1']}, # Standard format
95+
{'items': ['item2']}, # Alternative format
96+
{'results': ['item3']}, # Another format
97+
['item4'], # Direct list
98+
{'objects': ['item5']}, # Another common format
99+
]
100+
101+
call_count = 0
102+
def mock_fetch(params):
103+
nonlocal call_count
104+
call_count += 1
105+
if call_count <= len(responses):
106+
return responses[call_count - 1]
107+
return []
108+
109+
result = helper.paginate(mock_fetch)
110+
assert result == ['item1', 'item2', 'item3', 'item4', 'item5']
111+
112+
def test_paginate_empty_response(self) -> None:
113+
"""Test pagination with empty response."""
114+
helper = PaginationHelper()
115+
116+
def mock_fetch(params):
117+
return {'data': []}
118+
119+
result = helper.paginate(mock_fetch)
120+
assert result == []
121+
122+
def test_paginate_error_handling(self) -> None:
123+
"""Test pagination error handling."""
124+
helper = PaginationHelper(page_size=1)
125+
126+
call_count = 0
127+
def mock_fetch(params):
128+
nonlocal call_count
129+
call_count += 1
130+
if call_count == 1:
131+
return {'data': ['item1']}
132+
elif call_count == 2:
133+
raise ValueError("Page not found")
134+
return {'data': []}
135+
136+
# Should stop when error indicates end of pagination
137+
result = helper.paginate(mock_fetch)
138+
assert result == ['item1']
139+
assert call_count == 2
140+
141+
def test_paginate_response_object_with_attributes(self) -> None:
142+
"""Test pagination with response objects that have attributes."""
143+
helper = PaginationHelper(page_size=1)
144+
145+
# Mock response objects
146+
class MockResponse:
147+
def __init__(self, data):
148+
self.data = data
149+
150+
call_count = 0
151+
def mock_fetch(params):
152+
nonlocal call_count
153+
call_count += 1
154+
if call_count == 1:
155+
return MockResponse(['item1'])
156+
return MockResponse([])
157+
158+
result = helper.paginate(mock_fetch)
159+
assert result == ['item1']
160+
161+
def test_paginate_partial_page_stops_pagination(self) -> None:
162+
"""Test that getting fewer items than page_size stops pagination."""
163+
helper = PaginationHelper(page_size=3)
164+
165+
call_count = 0
166+
def mock_fetch(params):
167+
nonlocal call_count
168+
call_count += 1
169+
if call_count == 1:
170+
return {'data': ['item1', 'item2']} # Fewer than page_size
171+
return {'data': ['item3', 'item4', 'item5']} # This shouldn't be called
172+
173+
result = helper.paginate(mock_fetch)
174+
assert result == ['item1', 'item2']
175+
assert call_count == 1
176+
177+
@pytest.mark.asyncio
178+
async def test_paginate_async(self) -> None:
179+
"""Test async pagination functionality."""
180+
helper = PaginationHelper(page_size=2, max_pages=2)
181+
182+
call_count = 0
183+
async def mock_fetch(params):
184+
nonlocal call_count
185+
call_count += 1
186+
if call_count == 1:
187+
return {'data': ['item1', 'item2']}
188+
elif call_count == 2:
189+
return {'data': ['item3']} # Last page
190+
return {'data': []}
191+
192+
result = await helper.paginate_async(mock_fetch)
193+
assert result == ['item1', 'item2', 'item3']
194+
assert call_count == 2
195+
196+
def test_extract_items_edge_cases(self) -> None:
197+
"""Test _extract_items method with edge cases."""
198+
helper = PaginationHelper()
199+
200+
# Test None response
201+
assert helper._extract_items(None) == []
202+
203+
# Test dict without expected keys
204+
assert helper._extract_items({'other_key': 'value'}) == []
205+
206+
# Test dict with non-list values
207+
assert helper._extract_items({'data': 'not_a_list'}) == []
208+
209+
# Test nested structures
210+
assert helper._extract_items({'data': {'nested': ['item']}}) == []
211+
212+
def test_is_pagination_end_error(self) -> None:
213+
"""Test _is_pagination_end_error method."""
214+
helper = PaginationHelper()
215+
216+
# Should return True for pagination end errors
217+
assert helper._is_pagination_end_error(ValueError("Page not found"))
218+
assert helper._is_pagination_end_error(ValueError("Invalid page"))
219+
assert helper._is_pagination_end_error(ValueError("No more pages"))
220+
221+
# Should return False for other errors
222+
assert not helper._is_pagination_end_error(ValueError("Network error"))
223+
assert not helper._is_pagination_end_error(ValueError("Timeout"))
224+
225+
# Should handle different error types
226+
assert helper._is_pagination_end_error(Exception("page not found"))
227+
assert not helper._is_pagination_end_error(Exception("other error"))

0 commit comments

Comments
 (0)