feat(SYPL-5699): use get_page

Author: Daniel Bush

sypht/client.py

@@ -421,12 +421,11 @@ def get_annotations(
     ):
         page_iter = fetch_all_pages(
             name="get_annotations",
-            fetch_page=lambda *args, **kwargs: self._get_annotations(*args, **kwargs)[
-                "annotations"
-            ],
+            fetch_page=lambda *args, **kwargs: self._get_annotations(*args, **kwargs),
+            get_page=lambda response: response["annotations"],
         )
         annotations = []
-        for page in page_iter(
+        for response in page_iter(
             doc_id=doc_id,
             task_id=task_id,
             user_id=user_id,
@@ -435,7 +434,7 @@ def get_annotations(
             to_date=to_date,
             endpoint=endpoint,
         ):
-            annotations.extend(page)
+            annotations.extend(response["annotations"])
         return {"annotations": annotations}
 
     def _get_annotations(
@@ -477,14 +476,15 @@ def get_annotations_for_docs(self, doc_ids, endpoint=None):
             name="get_annotations_for_docs",
             fetch_page=lambda *args, **kwargs: self._get_annotations_for_docs(
                 *args, **kwargs
-            )["annotations"],
+            ),
+            get_page=lambda response: response["annotations"],
         )
         annotations = []
-        for page in page_iter(
+        for response in page_iter(
             doc_ids=doc_ids,
             endpoint=endpoint,
         ):
-            annotations.extend(page)
+            annotations.extend(response["annotations"])
         return {"annotations": annotations}
 
     def _get_annotations_for_docs(self, doc_ids, endpoint=None, offset=0):

sypht/util.py

@@ -2,9 +2,16 @@
 
 
 def fetch_all_pages(
-    name: str, fetch_page: Callable[..., List[Any]], rec_limit=20000
-) -> Callable[..., Iterator[List[Any]]]:
-    """Returns an iterator that calls fetch_page with an offset that we increment by the number of pages fetched.  Stop if page returns empty list."""
+    name: str,
+    fetch_page: Callable[..., Any],
+    get_page: Callable[..., List[Any]] = lambda x: x,
+    rec_limit=20000,
+) -> Callable[..., Iterator[Any]]:
+    """Returns an iterator that calls fetch_page with an offset that we increment by the number of pages fetched.  Stop if page returns empty list.
+
+    :param fetch_page: the api call that takes an offset (zero-based) that fetches a page of results
+    :param get_page: a function that extracts the page from the response
+    """
 
     def fetch_all_pages(*args, **kwargs) -> Iterator[List[Any]]:
         page_count = 0
@@ -17,7 +24,7 @@ def fetch_all_pages(*args, **kwargs) -> Iterator[List[Any]]:
                     f"fetch_all_pages({name}): fetched more than {rec_limit} items.  Consider using adding or adjusting a filter to reduce the number of items fetched."
                 )
             try:
-                result = fetch_page(
+                response = fetch_page(
                     *args,
                     **kwargs,
                     offset=page_count - 1,
@@ -26,9 +33,15 @@ def fetch_all_pages(*args, **kwargs) -> Iterator[List[Any]]:
                 raise Exception(
                     f"Failed fetching for {name} for offset={page_count - 1}"
                 ) from err
-            if not result:
+            try:
+                page = get_page(response)
+            except Exception as err:
+                raise Exception(
+                    f"get_page failed to extract page from response for {name} for offset={page_count - 1}"
+                ) from err
+            if len(page) == 0:
                 break
-            recs += len(result)
-            yield result
+            recs += len(page)
+            yield response
 
     return fetch_all_pages

tests/tests_util.py

@@ -25,6 +25,30 @@ def fetch_something(offset, pages=1):
     assert results == [0, 1, 2, 3, 4]
 
 
+def test_fetch_all_pages_can_fetch_one_page_with_get_page():
+    # arrange
+    page_size = 5
+
+    def fetch_something(offset, pages=1):
+        pages0 = pages - 1
+        if offset > pages0:
+            return {"results": []}
+        start = offset * page_size
+        page = range(start, start + page_size)
+        return {"results": list(page)}
+
+    # act
+    page_iter = fetch_all_pages(
+        name="test1", fetch_page=fetch_something, get_page=lambda resp: resp["results"]
+    )
+    results = []
+    for resp in page_iter(pages=1):
+        results += resp["results"]
+
+    # assert
+    assert results == [0, 1, 2, 3, 4]
+
+
 def test_fetch_all_pages_can_fetch_several_pages():
     # arrange
     page_size = 5