diff --git a/botocore/paginate.py b/botocore/paginate.py index 96c7c0b6fd..578601ddad 100644 --- a/botocore/paginate.py +++ b/botocore/paginate.py @@ -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.""" @@ -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: @@ -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] = [] @@ -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() diff --git a/tests/unit/test_paginate.py b/tests/unit/test_paginate.py index e73db7c93c..9d5ae6d2e5 100644 --- a/tests/unit/test_paginate.py +++ b/tests/unit/test_paginate.py @@ -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()