1# --------------------------------------------------------------------------
2#
3# Copyright (c) Microsoft Corporation. All rights reserved.
4#
5# The MIT License (MIT)
6#
7# Permission is hereby granted, free of charge, to any person obtaining a copy
8# of this software and associated documentation files (the ""Software""), to
9# deal in the Software without restriction, including without limitation the
10# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
11# sell copies of the Software, and to permit persons to whom the Software is
12# furnished to do so, subject to the following conditions:
13#
14# The above copyright notice and this permission notice shall be included in
15# all copies or substantial portions of the Software.
16#
17# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
22# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
23# IN THE SOFTWARE.
24#
25# --------------------------------------------------------------------------
26import itertools
27from typing import (
28 Callable,
29 Optional,
30 TypeVar,
31 Iterator,
32 Iterable,
33 Tuple,
34 Any,
35)
36import logging
37
38from .exceptions import AzureError
39
40
41_LOGGER = logging.getLogger(__name__)
42
43ReturnType = TypeVar("ReturnType")
44ResponseType = TypeVar("ResponseType")
45
46
47class PageIterator(Iterator[Iterator[ReturnType]]):
48 def __init__(
49 self,
50 get_next: Callable[[Optional[str]], ResponseType],
51 extract_data: Callable[[ResponseType], Tuple[str, Iterable[ReturnType]]],
52 continuation_token: Optional[str] = None,
53 ):
54 """Return an iterator of pages.
55
56 :param get_next: Callable that take the continuation token and return a HTTP response
57 :param extract_data: Callable that take an HTTP response and return a tuple continuation token,
58 list of ReturnType
59 :param str continuation_token: The continuation token needed by get_next
60 """
61 self._get_next = get_next
62 self._extract_data = extract_data
63 self.continuation_token = continuation_token
64 self._did_a_call_already = False
65 self._response: Optional[ResponseType] = None
66 self._current_page: Optional[Iterable[ReturnType]] = None
67
68 def __iter__(self) -> Iterator[Iterator[ReturnType]]:
69 return self
70
71 def __next__(self) -> Iterator[ReturnType]:
72 """Get the next page in the iterator.
73
74 :returns: An iterator of objects in the next page.
75 :rtype: iterator[ReturnType]
76 :raises StopIteration: If there are no more pages to return.
77 :raises AzureError: If the request fails.
78 """
79 if self.continuation_token is None and self._did_a_call_already:
80 raise StopIteration("End of paging")
81 try:
82 self._response = self._get_next(self.continuation_token)
83 except AzureError as error:
84 if not error.continuation_token:
85 error.continuation_token = self.continuation_token
86 raise
87
88 self._did_a_call_already = True
89
90 self.continuation_token, self._current_page = self._extract_data(self._response)
91
92 return iter(self._current_page)
93
94 next = __next__ # Python 2 compatibility. Can't be removed as some people are using ".next()" even in Py3
95
96
97class ItemPaged(Iterator[ReturnType]):
98 def __init__(self, *args: Any, **kwargs: Any) -> None:
99 """Return an iterator of items.
100
101 args and kwargs will be passed to the PageIterator constructor directly,
102 except page_iterator_class
103 """
104 self._args = args
105 self._kwargs = kwargs
106 self._page_iterator: Optional[Iterator[ReturnType]] = None
107 self._page_iterator_class = self._kwargs.pop("page_iterator_class", PageIterator)
108
109 def by_page(self, continuation_token: Optional[str] = None) -> Iterator[Iterator[ReturnType]]:
110 """Get an iterator of pages of objects, instead of an iterator of objects.
111
112 :param str continuation_token:
113 An opaque continuation token. This value can be retrieved from the
114 continuation_token field of a previous generator object. If specified,
115 this generator will begin returning results from this point.
116 :returns: An iterator of pages (themselves iterator of objects)
117 :rtype: iterator[iterator[ReturnType]]
118 """
119 return self._page_iterator_class(continuation_token=continuation_token, *self._args, **self._kwargs)
120
121 def __repr__(self) -> str:
122 return "<iterator object azure.core.paging.ItemPaged at {}>".format(hex(id(self)))
123
124 def __iter__(self) -> Iterator[ReturnType]:
125 return self
126
127 def __next__(self) -> ReturnType:
128 """Get the next item in the iterator.
129
130 :returns: The next item in the iterator.
131 :rtype: ReturnType
132 :raises StopIteration: If there are no more items to return.
133 """
134 if self._page_iterator is None:
135 self._page_iterator = itertools.chain.from_iterable(self.by_page())
136 return next(self._page_iterator)
137
138 next = __next__ # Python 2 compatibility. Can't be removed as some people are using ".next()" even in Py3