From 0970940bc507ae430fa312b4f390803fe3cba468 Mon Sep 17 00:00:00 2001
From: Raphael Krupinski <rafalkrupinski@users.noreply.github.com>
Date: Sun, 20 Oct 2024 21:56:01 +0200
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20`iter=5Fpages`=20to=20turn=20?=
 =?UTF-8?q?operation=20methods=20to=20async=20iterators.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ChangeLog.md                     |  1 +
 src/lapidary/runtime/__init__.py |  2 ++
 src/lapidary/runtime/paging.py   | 58 ++++++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+)
 create mode 100644 src/lapidary/runtime/paging.py

diff --git a/ChangeLog.md b/ChangeLog.md
index d796b71..3cbaab3 100644
--- a/ChangeLog.md
+++ b/ChangeLog.md
@@ -10,6 +10,7 @@ and the format of this file is based on [Keep a Changelog](https://keepachangelo
 
 ### Added
 - Accept session_factory in `ClientBase.__init__`.
+- Helper function to iterate over pages.
 
 ### Fixed
 - Handling collections in request bodies.
diff --git a/src/lapidary/runtime/__init__.py b/src/lapidary/runtime/__init__.py
index 14b8446..351597d 100644
--- a/src/lapidary/runtime/__init__.py
+++ b/src/lapidary/runtime/__init__.py
@@ -26,6 +26,7 @@
     'delete',
     'get',
     'head',
+    'iter_pages',
     'patch',
     'post',
     'put',
@@ -38,4 +39,5 @@
 from .model.error import HttpErrorResponse, LapidaryError, LapidaryResponseError, UnexpectedResponse
 from .model.param_serialization import Form, FormExplode, SimpleMultimap, SimpleString
 from .operation import delete, get, head, patch, post, put, trace
+from .paging import iter_pages
 from .types_ import ClientArgs, NamedAuth, SecurityRequirements, SessionFactory
diff --git a/src/lapidary/runtime/paging.py b/src/lapidary/runtime/paging.py
new file mode 100644
index 0000000..be71081
--- /dev/null
+++ b/src/lapidary/runtime/paging.py
@@ -0,0 +1,58 @@
+from collections.abc import AsyncIterable, Awaitable, Callable
+from typing import Optional, TypeVar
+
+from typing_extensions import ParamSpec, Unpack
+
+P = ParamSpec('P')
+R = TypeVar('R')
+C = TypeVar('C')
+
+
+def iter_pages(
+    fn: Callable[P, Awaitable[R]],
+    cursor_param_name: str,
+    get_cursor: Callable[[R], Optional[C]],
+) -> Callable[P, AsyncIterable[R]]:
+    """
+    Create a function that returns an async iterator over pages from the async operation function :param:`fn`.
+
+    The returned function can be called with the same parameters as :param:`fn` (except for the cursor parameter),
+    and returns an async iterator that yields results from :param:`fn`, handling pagination automatically.
+
+    The function :param:`fn` will be called initially without the cursor parameter and then called with the cursor parameter
+    as long as :param:`get_cursor` can extract a cursor from the result.
+
+    **Example:**
+
+    .. code:: python
+
+        async for page in iter_pages(client.fn, 'cursor', extractor_fn)(parameter=value):
+            # Process page
+
+    Typically, an API will use the same paging pattern for all operations supporting it, so it's a good idea to write a shortcut function:
+
+    .. code:: python
+
+        from lapidary.runtime import iter_pages as _iter_pages
+
+        def iter_pages[P, R](fn: Callable[P, Awaitable[R]]) -> Callable[P, AsyncIterable[R]]:
+            return _iter_pages(fn, 'cursor', lambda result: ...)
+
+    :param fn: An async function that retrieves a page of data.
+    :param cursor_param_name: The name of the cursor parameter in :param:`fn`.
+    :param get_cursor: A function that extracts a cursor value from the result of :param:`fn`. Return `None` to end the iteration.
+    """
+
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> AsyncIterable[R]:
+        result = await fn(*args, **kwargs)  # type: ignore[call-arg]
+        yield result
+        cursor = get_cursor(result)
+
+        while cursor:
+            kwargs[cursor_param_name] = cursor
+            result = await fn(*args, **kwargs)  # type: ignore[call-arg]
+            yield result
+
+            cursor = get_cursor(result)
+
+    return wrapper