First version of paging

sdk/core/azure-core/azure/core/async_paging.py

@@ -23,12 +23,14 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-from collections.abc import AsyncIterator
+from typing import AsyncIterator, TypeVar
 import logging
 
 _LOGGER = logging.getLogger(__name__)
 
-class AsyncPagedMixin(AsyncIterator):
+ReturnType = TypeVar("ReturnType")
+
+class AsyncPagedMixin(AsyncIterator[ReturnType]):
     """Bring async to Paging.
 
     **Keyword argument:**

sdk/core/azure-core/azure/core/paging.py

@@ -24,13 +24,8 @@
 #
 # --------------------------------------------------------------------------
 import sys
-try:
-    from collections.abc import Iterator
-    xrange = range
-except ImportError:
-    from collections import Iterator
 
-from typing import Dict, Any, List, Callable, Optional, TYPE_CHECKING  # pylint: disable=unused-import
+from typing import Dict, Any, List, Callable, Optional, TypeVar, Iterator, Iterable, Tuple, TYPE_CHECKING  # pylint: disable=unused-import
 
 if TYPE_CHECKING:
     from .pipeline.transport.base import HttpResponse
@@ -43,66 +38,74 @@
     class AsyncPagedMixin(object):  # type: ignore
         pass
 
-class Paged(AsyncPagedMixin, Iterator):
-    """A container for paged REST responses.
-
-    :param response: server response object.
-    :type response: ~azure.core.pipeline.transport.HttpResponse
-    :param callable command: Function to retrieve the next page of items.
-    :param Deserializer deserializer: a Deserializer instance to use
-    """
-    _validation = {}  # type: Dict[str, Dict[str, Any]]
-    _attribute_map = {}  # type: Dict[str, Dict[str, Any]]
-
-    def __init__(self, command, deserializer, **kwargs):
-        # type: (Callable[[str], HttpResponse], Deserializer, Any) -> None
-        super(Paged, self).__init__(**kwargs)  # type: ignore
-        # Sets next_link, current_page, and _current_page_iter_index.
-        self.next_link = ""
-        self.current_page = []  # type: List[Model]
-        self._current_page_iter_index = 0
-        self._deserializer = deserializer
-        self._get_next = command
-        self._response = None  # type: Optional[HttpResponse]
+ReturnType = TypeVar("ReturnType")
+
+
+class PageIterator(Iterator[Iterator[ReturnType]]):
+    def __init__(self, get_next, extract_data, continuation_token=None):
+        # type: (Callable[[str], ClientResponse], Callable[[ClientResponse], Tuple[str, List[ReturnType]], Optional[str]) -> None
+        """Return an iterator of pages.
+
+        :param get_next: Callable that take the continuation token and return a HTTP response
+        :param extract_data: Callable that take an HTTP response and return a tuple continuation token,
+         list of ReturnType
+        :param str continuation_token: The continuation token needed by get_next
+        """
+        self._get_next = get_next
+        self._extract_data = extract_data
+        self.continuation_token = continuation_token
+        self._did_a_call_already = False
+        self._current_page = []  # type: List[Model]
 
     def __iter__(self):
         """Return 'self'."""
-        # Since iteration mutates this object, consider it an iterator in-and-of
-        # itself.
         return self
 
-    @classmethod
-    def _get_subtype_map(cls):
-        """Required for parity to Model object for deserialization."""
-        return {}
+    def __next__(self):
+        if self.continuation_token is None and self._did_a_call_already:
+            raise StopIteration("End of paging")
 
-    def _advance_page(self):
-        # type: () -> List[Model]
-        """Force moving the cursor to the next azure call.
+        self._response = self._get_next(self.continuation_token)
+        self._did_a_call_already = True
 
-        This method is for advanced usage, iterator protocol is prefered.
+        self.continuation_token, self._current_page = self._extract_data(self._response)
+        return self._current_page
 
-        :raises: StopIteration if no further page
-        :return: The current page list
-        :rtype: list
-        """
-        if self.next_link is None:
-            raise StopIteration("End of paging")
-        self._current_page_iter_index = 0
-        self._response = self._get_next(self.next_link)
-        self._deserializer(self, self._response)
-        return self.current_page
+    next = __next__  # Python 2 compatibility.
+
+
+class ItemPaged(Iterable[ReturnType]):
+    def __init__(self, get_next, extract_data):
+        # type: (Callable[[str], ClientResponse], Callable[[ClientResponse], Tuple[str, List[ReturnType]]) -> None
+        self._get_next = get_next
+        self._extract_data = extract_data
+        self._page_iterator = None
+        self._page = None
+
+    def by_page(self, continuation_token=None):
+        # type: () -> PageIterator[ReturnType]
+        return PageIterator(
+            get_next=self._get_next,
+            extract_data=self._extract_data,
+            continuation_token=continuation_token
+        )
+
+    def __iter__(self):
+        """Return 'self'."""
+        return self
 
     def __next__(self):
-        """Iterate through responses."""
-        # Storing the list iterator might work out better, but there's no
-        # guarantee that some code won't replace the list entirely with a copy,
-        # invalidating an list iterator that might be saved between iterations.
-        if self.current_page and self._current_page_iter_index < len(self.current_page):
-            response = self.current_page[self._current_page_iter_index]
-            self._current_page_iter_index += 1
-            return response
-        self._advance_page()
-        return self.__next__()
+        if self._page_iterator is None:
+            self._page_iterator = self.by_page()
+            return next(self)
+        if self._page is None:
+            # Let it raise StopIteration
+            self._page = iter(next(self._page_iterator))
+            return next(self)
+        try:
+            return next(self._page)
+        except StopIteration:
+            self._page = None
+            return next(self)
 
     next = __next__  # Python 2 compatibility.

sdk/core/azure-core/tests/test_paging.py

@@ -1,6 +1,6 @@
 #--------------------------------------------------------------------------
 #
-# Copyright (c) Microsoft Corporation. All rights reserved. 
+# Copyright (c) Microsoft Corporation. All rights reserved.
 #
 # The MIT License (MIT)
 #
@@ -25,28 +25,19 @@
 #--------------------------------------------------------------------------
 
 import unittest
+import pytest
 
-from azure.core.paging import Paged
+from azure.core.paging import ItemPaged
 
-from msrest.serialization import Deserializer
-
-class FakePaged(Paged):
-    _attribute_map = {
-        'next_link': {'key': 'nextLink', 'type': 'str'},
-        'current_page': {'key': 'value', 'type': '[str]'}
-    }
-
-    def __init__(self, *args, **kwargs):
-        super(FakePaged, self).__init__(*args, **kwargs)
-
-_test_deserializer = Deserializer({})
 
 class TestPaging(unittest.TestCase):
 
     def test_basic_paging(self):
 
-        def internal_paging(next_link=None, raw=False):
-            if not next_link:
+        def get_next(continuation_token=None):
+            """Simplify my life and return JSON and not response, but should be response.
+            """
+            if not continuation_token:
                 return {
                     'nextLink': 'page2',
                     'value': ['value1.0', 'value1.1']
@@ -57,17 +48,51 @@ def internal_paging(next_link=None, raw=False):
                     'value': ['value2.0', 'value2.1']
                 }
 
-        deserialized = FakePaged(internal_paging, _test_deserializer)
-        result_iterated = list(deserialized)
+        def extract_data(response):
+            return response['nextLink'], response['value']
+
+        pager = ItemPaged(get_next, extract_data)
+        result_iterated = list(pager)
         self.assertListEqual(
             ['value1.0', 'value1.1', 'value2.0', 'value2.1'],
             result_iterated
         )
 
+    def test_by_page_paging(self):
+
+        def get_next(continuation_token=None):
+            """Simplify my life and return JSON and not response, but should be response.
+            """
+            if not continuation_token:
+                return {
+                    'nextLink': 'page2',
+                    'value': ['value1.0', 'value1.1']
+                }
+            else:
+                return {
+                    'nextLink': None,
+                    'value': ['value2.0', 'value2.1']
+                }
+
+        def extract_data(response):
+            return response['nextLink'], response['value']
+
+        pager = ItemPaged(get_next, extract_data).by_page()
+        page1 = next(pager)
+        assert list(page1) == ['value1.0', 'value1.1']
+
+        page2 = next(pager)
+        assert list(page2) == ['value2.0', 'value2.1']
+
+        with pytest.raises(StopIteration):
+            next(pager)
+
     def test_advance_paging(self):
 
-        def internal_paging(next_link=None, raw=False):
-            if not next_link:
+        def get_next(continuation_token=None):
+            """Simplify my life and return JSON and not response, but should be response.
+            """
+            if not continuation_token:
                 return {
                     'nextLink': 'page2',
                     'value': ['value1.0', 'value1.1']
@@ -78,27 +103,32 @@ def internal_paging(next_link=None, raw=False):
                     'value': ['value2.0', 'value2.1']
                 }
 
-        deserialized = FakePaged(internal_paging, _test_deserializer)
-        page1 = next(deserialized)
+        def extract_data(response):
+            return response['nextLink'], response['value']
+
+        pager = ItemPaged(get_next, extract_data)
+        page1 = next(pager)
         assert page1 == 'value1.0'
-        page1 = next(deserialized)
+        page1 = next(pager)
         assert page1 == 'value1.1'
-     
-        page2 = next(deserialized)
+
+        page2 = next(pager)
         assert page2 == 'value2.0'
-        page2 = next(deserialized)
+        page2 = next(pager)
         assert page2 == 'value2.1'
 
         with self.assertRaises(StopIteration):
-            next(deserialized)
+            next(pager)
 
     def test_none_value(self):
-        def internal_paging(next_link=None, raw=False):
+        def get_next(next_link=None):
             return {
                 'nextLink': None,
                 'value': None
             }
+        def extract_data(response):
+            return response['nextLink'], response['value'] or []
 
-        deserialized = FakePaged(internal_paging, _test_deserializer)
-        result_iterated = list(deserialized)
+        pager = ItemPaged(get_next, extract_data)
+        result_iterated = list(pager)
         self.assertEqual(len(result_iterated), 0)