Skip to content

Commit

Permalink
Expose result_keys on PageIterator object
Browse files Browse the repository at this point in the history
This was expose on the Paginator object, but also needs to be
exposed through the PageIterator object.

Fixes #442.
  • Loading branch information
jamesls committed Oct 28, 2013
1 parent ed5b0b4 commit 0a98edb
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 6 deletions.
16 changes: 10 additions & 6 deletions botocore/paginate.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,18 +105,22 @@ def _extract_paging_params(self, kwargs):

class PageIterator(object):
def __init__(self, operation, input_token, output_token, more_results,
result_key, max_items, starting_token, endpoint, op_kwargs):
result_keys, max_items, starting_token, endpoint, op_kwargs):
self._operation = operation
self._input_token = input_token
self._output_token = output_token
self._more_results = more_results
self._result_key = result_key
self._result_keys = result_keys
self._max_items = max_items
self._starting_token = starting_token
self._endpoint = endpoint
self._op_kwargs = op_kwargs
self._resume_token = None

@property
def result_keys(self):
return self._result_keys

@property
def resume_token(self):
"""Token to specify to resume pagination."""
Expand All @@ -135,7 +139,7 @@ def __iter__(self):
# The number of items from result_key we've seen so far.
total_items = 0
first_request = True
primary_result_key = self._result_key[0]
primary_result_key = self.result_keys[0]
starting_truncation = 0
self._inject_starting_token(current_kwargs)
while True:
Expand Down Expand Up @@ -205,7 +209,7 @@ def _handle_first_request(self, parsed, primary_result_key,
# We also need to truncate any secondary result keys
# because they were not truncated in the previous last
# response.
for token in self._result_key:
for token in self.result_keys:
if token == primary_result_key:
continue
parsed[token] = []
Expand Down Expand Up @@ -241,9 +245,9 @@ def _get_next_token(self, parsed):
return next_tokens

def result_key_iters(self):
teed_results = tee(self, len(self._result_key))
teed_results = tee(self, len(self.result_keys))
return [ResultKeyIterator(i, result_key) for i, result_key
in zip(teed_results, self._result_key)]
in zip(teed_results, self.result_keys)]

def build_full_result(self):
iterators = self.result_key_iters()
Expand Down
7 changes: 7 additions & 0 deletions tests/unit/test_paginate.py
Original file line number Diff line number Diff line change
Expand Up @@ -514,6 +514,13 @@ def test_resume_with_multiple_input_keys(self):
self.operation.call.call_args_list,
[mock.call(None, InMarker1='m1', InMarker2='m2'),])

def test_result_key_exposed_on_paginator(self):
self.assertEqual(self.paginator.result_keys, ['Users', 'Groups'])

def test_result_key_exposed_on_page_iterator(self):
pages = self.paginator.paginate(None, max_items=3)
self.assertEqual(pages.result_keys, ['Users', 'Groups'])


if __name__ == '__main__':
unittest.main()

0 comments on commit 0a98edb

Please sign in to comment.