Merge pull request #193 from ntropy-network/auto-paginate

Add PagedResponse and auto_paginate

.github/workflows/python-app.yml

@@ -22,10 +22,10 @@ jobs:
     - name: Mask sensitive
       run: |
           echo "::add-mask::${{ secrets.NTROPY_API_KEY }}"
-    - name: Set up Python 3.6
+    - name: Set up Python 3.8
       uses: actions/setup-python@v2
       with:
-        python-version: 3.6.15
+        python-version: 3.8.18
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

ntropy_sdk/v3/__init__.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING, Optional
 import requests
 
+
 if TYPE_CHECKING:
     from ntropy_sdk.ntropy_sdk import SDK
     from typing_extensions import TypedDict

ntropy_sdk/v3/account_holders.py

@@ -1,11 +1,12 @@
 from datetime import datetime
 from enum import Enum
-from typing import TYPE_CHECKING, List, Optional
+from typing import Optional, TYPE_CHECKING
 import uuid
 
 from pydantic import BaseModel, Field
 
 from ntropy_sdk.utils import pydantic_json
+from ntropy_sdk.v3.paging import PagedResponse
 
 if TYPE_CHECKING:
     from ntropy_sdk.ntropy_sdk import SDK
@@ -45,7 +46,7 @@ def list(
         cursor: Optional[str] = None,
         limit: Optional[int] = None,
         **extra_kwargs: "Unpack[ExtraKwargs]",
-    ) -> List[AccountHolder]:
+    ) -> PagedResponse[AccountHolder]:
         """List all account holders"""
 
         request_id = extra_kwargs.get("request_id")
@@ -63,7 +64,15 @@ def list(
             },
             **extra_kwargs,
         )
-        return [AccountHolder(**j, request_id=request_id) for j in resp.json()["data"]]
+        page = PagedResponse[AccountHolder](
+            **resp.json(),
+            request_id=request_id,
+            _resource=self,
+            _extra_kwargs=extra_kwargs,
+        )
+        for t in page.data:
+            t.request_id = request_id
+        return page
 
     def get(self, *, id: str, **extra_kwargs: "Unpack[ExtraKwargs]") -> AccountHolder:
         """Retrieve an account holder"""

ntropy_sdk/v3/bank_statements.py

@@ -2,14 +2,15 @@
 from enum import Enum
 from io import IOBase
 import time
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import List, Optional, TYPE_CHECKING, Union
 import uuid
 
 from pydantic import BaseModel, Field, NonNegativeFloat
 
 from ntropy_sdk.bank_statements import StatementInfo
 from ntropy_sdk.errors import NtropyDatasourceError
 from ntropy_sdk.utils import EntryType
+from ntropy_sdk.v3.paging import PagedResponse
 
 if TYPE_CHECKING:
     from ntropy_sdk.ntropy_sdk import SDK
@@ -115,7 +116,7 @@ def list(
         limit: Optional[int] = None,
         status: Optional[BankStatementJobStatus] = None,
         **extra_kwargs: "Unpack[ExtraKwargs]",
-    ) -> List[BankStatementJob]:
+    ) -> PagedResponse[BankStatementJob]:
         request_id = extra_kwargs.get("request_id")
         if request_id is None:
             request_id = uuid.uuid4().hex
@@ -133,9 +134,15 @@ def list(
             payload=None,
             **extra_kwargs,
         )
-        return [
-            BankStatementJob(**j, request_id=request_id) for j in resp.json()["data"]
-        ]
+        page = PagedResponse[BankStatementJob](
+            **resp.json(),
+            request_id=request_id,
+            _resource=self,
+            _extra_kwargs=extra_kwargs,
+        )
+        for b in page.data:
+            b.request_id = request_id
+        return page
 
     def create(
         self,

ntropy_sdk/v3/batches.py

@@ -1,13 +1,14 @@
 from datetime import datetime
 from enum import Enum
 import time
-from typing import TYPE_CHECKING, List, Optional
+from typing import List, Optional, TYPE_CHECKING
 import uuid
 
 from pydantic import BaseModel, Field
 
 from ntropy_sdk.errors import NtropyBatchError
 from ntropy_sdk.utils import pydantic_json
+from ntropy_sdk.v3.paging import PagedResponse
 from ntropy_sdk.v3.transactions import EnrichedTransaction, EnrichmentInput
 
 if TYPE_CHECKING:
@@ -95,7 +96,7 @@ def list(
         limit: Optional[int] = None,
         status: Optional[BatchStatus] = None,
         **extra_kwargs: "Unpack[ExtraKwargs]",
-    ) -> List[Batch]:
+    ) -> PagedResponse[Batch]:
         """List all batches"""
 
         request_id = extra_kwargs.get("request_id")
@@ -114,7 +115,15 @@ def list(
             },
             **extra_kwargs,
         )
-        return [Batch(**j, request_id=request_id) for j in resp.json()["data"]]
+        page = PagedResponse[Batch](
+            **resp.json(),
+            request_id=request_id,
+            _resource=self,
+            _extra_kwargs=extra_kwargs,
+        )
+        for t in page.data:
+            t.request_id = request_id
+        return page
 
     def get(self, *, id: str, **extra_kwargs: "Unpack[ExtraKwargs]") -> Batch:
         """Retrieve a batch"""

ntropy_sdk/v3/paging.py

@@ -0,0 +1,117 @@
+from dataclasses import dataclass
+from typing import (
+    Any,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Protocol,
+    TYPE_CHECKING,
+    TypeVar,
+)
+
+from pydantic import PrivateAttr
+
+from ntropy_sdk.utils import PYDANTIC_V2
+
+if TYPE_CHECKING:
+    from . import ExtraKwargs
+    from pydantic import BaseModel as GenericModel
+    from typing_extensions import Unpack, Self
+
+elif PYDANTIC_V2:
+    from pydantic import BaseModel as GenericModel
+else:
+    from pydantic.generics import GenericModel
+
+T = TypeVar("T")
+
+
+class ListableResource(Protocol[T]):
+    def list(
+        self,
+        *,
+        cursor: str,
+        limit: Optional[int],
+        **extra_kwargs: "Unpack[ExtraKwargs]",
+    ) -> "PagedResponse[T]":
+        ...
+
+
+class PagedResponse(GenericModel, Generic[T]):
+    next_cursor: Optional[str]
+    data: List[T]
+    request_id: Optional[str] = None
+    _resource: Optional[ListableResource[T]] = PrivateAttr(None)
+    if TYPE_CHECKING:
+        _extra_kwargs: Optional["ExtraKwargs"] = PrivateAttr(None)
+        pass
+    else:
+        _extra_kwargs: Any = PrivateAttr()  # pydantic v1 complains about ExtraKwargs
+
+        def __init__(
+            self,
+            *,
+            _resource: Optional[ListableResource[T]] = None,
+            _extra_kwargs: Optional["ExtraKwargs"] = None,
+            **data,
+        ):
+            super().__init__(**data)
+            self._resource = _resource
+            self._extra_kwargs = _extra_kwargs
+
+    def auto_paginate(
+        self,
+        *,
+        page_size: Optional[int] = None,
+    ) -> "AutoPaginate[T]":
+        if self._resource is None:
+            raise ValueError("self._resource is None")
+        return AutoPaginate(
+            _first_page=self,
+            _resource=self._resource,
+            _page_size=page_size,
+        )
+
+
+@dataclass
+class AutoPaginate(Generic[T]):
+    _first_page: PagedResponse
+    _resource: ListableResource[T]
+    _page_size: Optional[int]
+
+    def __iter__(self) -> "AutoPaginateIterator[T]":
+        return AutoPaginateIterator(
+            current_iter=iter(self._first_page.data),
+            page_size=self._page_size,
+            next_cursor=self._first_page.next_cursor,
+            _resource=self._resource,
+            _extra_kwargs=self._first_page._extra_kwargs or {},
+        )
+
+
+@dataclass
+class AutoPaginateIterator(Generic[T]):
+    current_iter: Iterator[T]
+    next_cursor: Optional[str]
+    page_size: Optional[int]
+    _resource: ListableResource[T]
+    _extra_kwargs: "ExtraKwargs"
+
+    def __next__(self) -> T:
+        try:
+            return next(self.current_iter)
+        except StopIteration:
+            if self.next_cursor is None:
+                raise StopIteration
+            next_page = self._resource.list(
+                cursor=self.next_cursor,
+                limit=self.page_size,
+                **self._extra_kwargs,
+            )
+            self.current_iter = iter(next_page.data)
+            self.next_cursor = next_page.next_cursor
+            return next(self.current_iter)
+
+    def __iter__(self) -> "Self":
+        return self

ntropy_sdk/v3/transactions.py

@@ -1,9 +1,12 @@
 from datetime import date as dt_date, date, datetime
 import enum
-from typing import TYPE_CHECKING, List, Optional
+from typing import List, Optional, TYPE_CHECKING
 import uuid
+
 from pydantic import BaseModel, Field, NonNegativeFloat
-from ntropy_sdk.utils import PYDANTIC_V2, EntryType, pydantic_json
+
+from ntropy_sdk.utils import EntryType, PYDANTIC_V2, pydantic_json
+from ntropy_sdk.v3.paging import PagedResponse
 
 PYDANTIC_PATTERN = "pattern" if PYDANTIC_V2 else "regex"
 MAX_SYNC_BATCH = 1000
@@ -310,7 +313,7 @@ def list(
         account_holder_id: Optional[str] = None,
         dataset_id: Optional[int] = None,
         **extra_kwargs: "Unpack[ExtraKwargs]",
-    ) -> List[Transaction]:
+    ) -> PagedResponse[Transaction]:
         """List all transactions"""
 
         request_id = extra_kwargs.get("request_id")
@@ -331,7 +334,15 @@ def list(
             payload=None,
             **extra_kwargs,
         )
-        return [Transaction(**j, request_id=request_id) for j in resp.json()["data"]]
+        page = PagedResponse[Transaction](
+            **resp.json(),
+            request_id=request_id,
+            _resource=self,
+            _extra_kwargs=extra_kwargs,
+        )
+        for t in page.data:
+            t.request_id = request_id
+        return page
 
     def get(self, *, id: str, **extra_kwargs: "Unpack[ExtraKwargs]") -> Transaction:
         """Retrieve a transaction"""

setup.py

@@ -26,14 +26,14 @@
 setup(
     author="Ntropy",
     author_email="[email protected]",
-    python_requires=">=3.6",
+    python_requires=">=3.8",
     classifiers=[
         "Development Status :: 2 - Pre-Alpha",
         "Intended Audience :: Developers",
         "License :: OSI Approved :: MIT License",
         "Natural Language :: English",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.6",
+        "Programming Language :: Python :: 3.8",
     ],
     description="SDK for the Ntropy API",
     entry_points={

tests/test_v3.py

@@ -0,0 +1,10 @@
+from itertools import islice
+from ntropy_sdk import SDK
+
+
+def test_pagination(sdk: SDK):
+    tx_ids = set()
+    it = sdk.v3.transactions.list(limit=2).auto_paginate(page_size=2)
+    for tx in islice(it, 10):
+        tx_ids.add(tx.id)
+    assert len(tx_ids) == 10