Add a generator to iterate through paginated methods

Author: rossjrw

scuttle/versions/v1.py

@@ -18,6 +18,19 @@ def wrapper(method, instance, args, kwargs):
         return instance.request(endpoint_url, *method(*args, **kwargs))
     return wrapper
 
+def get_default_data(**kwargs):
+    """Returns a POST data dict using default values."""
+    limit = kwargs.get('limit', 20)
+    offset = kwargs.get('offset', 0)
+    direction = kwargs.get('direction', 'asc')
+    if not isinstance(limit, int):
+        raise TypeError("`limit` must be int")
+    if not isinstance(offset, int):
+        raise TypeError("`offset` must be int")
+    if not direction in ('asc', 'desc'):
+        raise ValueError("`direction` must be one of 'asc', 'desc'")
+    return { 'limit': limit, 'offset': offset, 'direction': direction }
+
 class PaginatedMethod:
     """Object representing a method that has a POST paginated version (with
     args like limit, offset, direction) as well as optionally a GET
@@ -38,12 +51,8 @@ def __call__(self, *args, **kwargs):
             )
         return self._method(*args, **kwargs)
 
-    def verbose(self, *args, limit=None, offset=None, direction=None):
-        data = {
-            'limit': 20 if limit is None else limit,
-            'offset': 0 if offset is None else offset,
-            'direction': 'asc' if direction is None else direction,
-        }
+    def verbose(self, *args, **kwargs):
+        data = get_default_data(**kwargs)
         return self.__call__(*args, data=data, __verbose=True)
 
 class Api(BaseApi):
@@ -64,6 +73,49 @@ def __init__(self, *args):
         self.wikidotuser_revisions = PaginatedMethod(self._wikidotuser_revisions)
         self.tags_pages = PaginatedMethod(self._tags_pages, True)
 
+    def verbose(self, method: Callable, *args, **kwargs):
+        """Returns a generator that iterates over a paginated method.
+
+        callable `method`: The paginated method to iterate.
+        Remaining arguments will be passed to this method.
+
+
+        Pass this function int `limit`, an initial int `offset`, and str
+        `direction`. Each time the returned generator is called, it will
+        increment `offset` by `limit` and return the method for the resulting
+        set of parameters. Effectively, applied to a paginated method, this
+        generator is the same as turning the page.
+
+            wiki = scuttle(domain, API_KEY, 1)
+            generator = wiki.verbose(wiki.thread_posts, thread_id, limit=5)
+            for posts in generator:
+                print(len(posts))  # will be 5, except at very end
+
+        Note that at the end of the data, the length of the final 'page' will
+        very likely be less than `limit`.
+        """
+        print("hello!")
+        print(method, isinstance(method, PaginatedMethod))
+        if not isinstance(method, PaginatedMethod):
+            raise TypeError("Iterated method must be a paginated method")
+        data = get_default_data(**kwargs)
+
+        def paginated_generator():
+            limit: int = data['limit']
+            length: int = limit
+            offset: int = data['offset']
+            direction: str = data['direction']
+            while True:
+                result = method.verbose(*args, limit=limit, offset=offset,
+                                direction=direction)
+                yield result
+                if length < data['limit']:
+                    return
+                length = len(result)
+                offset += length
+
+        return paginated_generator()
+
     @endpoint("wikis")
     def wikis(self):
         return None, None

test/test_api.py

@@ -35,18 +35,31 @@ def test_pagination():
     page_slug = "main"
     page_id = wiki.page_by_slug(page_slug)['id']
     # non-paginated revisions - should just be metadata
-    print("DOING non-paginated")
     non_paginated_revisions = wiki.page_revisions(page_id)
-    print("DONE non-paginated")
     assert len(non_paginated_revisions) > 100
     assert 'content' not in non_paginated_revisions[0].keys()
     # paginated revisions - should include revision content
-    print("DOING paginated")
     paginated_revisions = wiki.page_revisions.verbose(page_id)
-    print("DONE paginated")
     assert 'content' in paginated_revisions[0].keys()
     assert len(paginated_revisions) == 20
 
+def test_pagination_generator():
+    wiki = scuttle.scuttle('en', API_KEY, 1)
+    # make a generator
+    page_slug = "main"
+    page_id = wiki.page_by_slug(page_slug)['id']
+    gen1 = wiki.verbose(wiki.page_revisions, page_id, limit=100)
+    assert int(next(gen1)[0]['metadata']['wikidot_metadata']['revision_number']) == 0
+    assert int(next(gen1)[0]['metadata']['wikidot_metadata']['revision_number']) == 100
+    # make another generator, see if they interfere
+    gen2 = wiki.verbose(wiki.page_revisions, page_id, limit=10, offset=10)
+    assert int(next(gen2)[0]['metadata']['wikidot_metadata']['revision_number']) == 10
+    assert int(next(gen2)[0]['metadata']['wikidot_metadata']['revision_number']) == 20
+    assert int(next(gen1)[0]['metadata']['wikidot_metadata']['revision_number']) == 200
+    # check errors
+    with pytest.raises(TypeError):
+        wiki.verbose(len)
+
 def test_page():
     wiki = scuttle.scuttle('en', API_KEY, 1)
     pages = wiki.pages()