diff --git a/flareio/api_client.py b/flareio/api_client.py index 7c3f0aa..a6ec673 100644 --- a/flareio/api_client.py +++ b/flareio/api_client.py @@ -146,3 +146,54 @@ def delete( json=json, headers=headers, ) + + def scroll( + self, + *, + method: t.Literal[ + "GET", + "POST", + ], + url: str, + params: t.Optional[t.Dict[str, t.Any]] = None, + json: t.Optional[t.Dict[str, t.Any]] = None, + headers: t.Optional[t.Dict[str, t.Any]] = None, + ) -> t.Iterator[requests.Response]: + if method not in {"GET", "POST"}: + raise Exception("Scrolling is only supported for GET or POST") + + from_in_params: bool = "from" in (params or {}) + from_in_json: bool = "from" in (json or {}) + + if not (from_in_params or from_in_json): + raise Exception("You must specify from either in params or in json") + if from_in_params and from_in_json: + raise Exception("You can't specify from both in params and in json") + + while True: + resp = self._request( + method=method, + url=url, + params=params, + json=json, + headers=headers, + ) + resp.raise_for_status() + + yield resp + + resp_body = resp.json() + + if "next" not in resp_body: + raise Exception( + "'next' was not found in the response body. Are you sure it supports scrolling?" + ) + + next_page: t.Optional[str] = resp_body["next"] + if not next_page: + break + + if params and from_in_params: + params["from"] = next_page + if json and from_in_json: + json["from"] = next_page diff --git a/tests/test_api_client.py b/tests/test_api_client.py index f26cbce..5d5ee1e 100644 --- a/tests/test_api_client.py +++ b/tests/test_api_client.py @@ -1,4 +1,5 @@ import pytest +import requests import requests_mock from datetime import datetime @@ -163,3 +164,76 @@ def test_wrapped_methods() -> None: client.delete("https://api.flare.io/hello-delete") assert mocker.last_request.url == "https://api.flare.io/hello-delete" assert mocker.last_request.headers["Authorization"] == "Bearer test-token-hello" + + +def test_scroll() -> None: + api_client = _get_test_client() + + # This should make no http call. + with requests_mock.Mocker() as mocker: + response_iterator: t.Iterator[requests.Response] = api_client.scroll( + method="GET", + url="https://api.flare.io/leaksdb/v2/sources", + params={ + "from": None, + }, + ) + assert len(mocker.request_history) == 0 + + # First page + with requests_mock.Mocker() as mocker: + mocker.register_uri( + "GET", + "https://api.flare.io/leaksdb/v2/sources", + json={ + "items": [1, 2, 3], + "next": "second_page", + }, + status_code=200, + ) + resp: requests.Response = next(response_iterator) + assert resp.json() == {"items": [1, 2, 3], "next": "second_page"} + + assert len(mocker.request_history) == 1 + assert mocker.last_request.query == "" + + # Second Page + with requests_mock.Mocker() as mocker: + mocker.register_uri( + "GET", + "https://api.flare.io/leaksdb/v2/sources?from=second_page", + json={ + "items": [4, 5, 6], + "next": "third_page", + }, + status_code=200, + ) + resp = next(response_iterator) + assert resp.json() == {"items": [4, 5, 6], "next": "third_page"} + assert len(mocker.request_history) == 1 + assert mocker.last_request.query == "from=second_page" + + # Third Page + with requests_mock.Mocker() as mocker: + mocker.register_uri( + "GET", + "https://api.flare.io/leaksdb/v2/sources?from=third_page", + json={ + "items": [7, 8, 9], + "next": None, + }, + status_code=200, + ) + resp = next(response_iterator) + assert resp.json() == { + "items": [7, 8, 9], + "next": None, + } + assert len(mocker.request_history) == 1 + assert mocker.last_request.query == "from=third_page" + + # We stop here. + with requests_mock.Mocker() as mocker: + with pytest.raises(StopIteration): + resp = next(response_iterator) + assert len(mocker.request_history) == 0