Skip to content

Commit

Permalink
Update load_all_pages.py
Browse files Browse the repository at this point in the history
  • Loading branch information
saleweaver authored Dec 20, 2024
1 parent 45960d6 commit 273b20f
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion sp_api/util/load_all_pages.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ def load_all_pages(
next_token_param="NextToken",
use_rate_limit_header: bool = False,
extras: dict = None,
next_token_only: bool = False
):
"""
Load all pages if a next token is returned
Expand All @@ -21,6 +22,7 @@ def load_all_pages(
next_token_param: str | The param amazon expects to hold the next token
use_rate_limit_header: if the function should try to use amazon's rate limit header
extras: additional data to be sent with NextToken, e.g `dict(QueryType='NEXT_TOKEN')` for `FulfillmentInbound`
next_token_only: remove all other params from kwargs, required for reports API
Returns:
Transforms the function in a generator, returning all pages
"""
Expand All @@ -39,7 +41,10 @@ def wrapper(*args, **kwargs):
)
if sleep_time > 0:
time.sleep(sleep_time)
kwargs.update({next_token_param: res.next_token, **extras})
if next_token_only:
kwargs = {next_token_param: res.next_token}
else:
kwargs.update({next_token_param: res.next_token, **extras})
else:
done = True

Expand Down

0 comments on commit 273b20f

Please sign in to comment.