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 (  # pylint: disable=unused-import
28    Callable,
29    Optional,
30    TypeVar,
31    Iterator,
32    Iterable,
33    Tuple,
34)
35import logging
36
37from .exceptions import AzureError
38
39
40_LOGGER = logging.getLogger(__name__)
41
42ReturnType = TypeVar("ReturnType")
43ResponseType = TypeVar("ResponseType")
44
45
46class PageIterator(Iterator[Iterator[ReturnType]]):
47    def __init__(
48        self,
49        get_next,  # type: Callable[[Optional[str]], ResponseType]
50        extract_data,  # type: Callable[[ResponseType], Tuple[str, Iterable[ReturnType]]]
51        continuation_token=None,  # type: Optional[str]
52    ):
53        """Return an iterator of pages.
54
55        :param get_next: Callable that take the continuation token and return a HTTP response
56        :param extract_data: Callable that take an HTTP response and return a tuple continuation token,
57         list of ReturnType
58        :param str continuation_token: The continuation token needed by get_next
59        """
60        self._get_next = get_next
61        self._extract_data = extract_data
62        self.continuation_token = continuation_token
63        self._did_a_call_already = False
64        self._response = None  # type: Optional[ResponseType]
65        self._current_page = None  # type: Optional[Iterable[ReturnType]]
66
67    def __iter__(self):
68        """Return 'self'."""
69        return self
70
71    def __next__(self):
72        # type: () -> Iterator[ReturnType]
73        if self.continuation_token is None and self._did_a_call_already:
74            raise StopIteration("End of paging")
75        try:
76            self._response = self._get_next(self.continuation_token)
77        except AzureError as error:
78            if not error.continuation_token:
79                error.continuation_token = self.continuation_token
80            raise
81
82        self._did_a_call_already = True
83
84        self.continuation_token, self._current_page = self._extract_data(self._response)
85
86        return iter(self._current_page)
87
88    next = __next__  # Python 2 compatibility.
89
90
91class ItemPaged(Iterator[ReturnType]):
92    def __init__(self, *args, **kwargs):
93        """Return an iterator of items.
94
95        args and kwargs will be passed to the PageIterator constructor directly,
96        except page_iterator_class
97        """
98        self._args = args
99        self._kwargs = kwargs
100        self._page_iterator = None
101        self._page_iterator_class = self._kwargs.pop(
102            "page_iterator_class", PageIterator
103        )
104
105    def by_page(self, continuation_token=None):
106        # type: (Optional[str]) -> Iterator[Iterator[ReturnType]]
107        """Get an iterator of pages of objects, instead of an iterator of objects.
108
109        :param str continuation_token:
110            An opaque continuation token. This value can be retrieved from the
111            continuation_token field of a previous generator object. If specified,
112            this generator will begin returning results from this point.
113        :returns: An iterator of pages (themselves iterator of objects)
114        """
115        return self._page_iterator_class(
116            continuation_token=continuation_token, *self._args, **self._kwargs
117        )
118
119    def __repr__(self):
120        return "<iterator object azure.core.paging.ItemPaged at {}>".format(hex(id(self)))
121
122    def __iter__(self):
123        """Return 'self'."""
124        return self
125
126    def __next__(self):
127        if self._page_iterator is None:
128            self._page_iterator = itertools.chain.from_iterable(self.by_page())
129        return next(self._page_iterator)
130
131    next = __next__  # Python 2 compatibility.
132