diff --git a/flareio/_ratelimit.py b/flareio/_ratelimit.py new file mode 100644 index 0000000..91fe94c --- /dev/null +++ b/flareio/_ratelimit.py @@ -0,0 +1,43 @@ +import time + +from datetime import datetime +from datetime import timedelta + +import typing as t + + +class _Limiter: + def __init__( + self, + *, + tick_interval: timedelta, + _sleeper: t.Callable[[float], None] = time.sleep, + ) -> None: + self._tick_interval: timedelta = tick_interval + self._next_tick: datetime = datetime.now() + self._sleeper: t.Callable[[float], None] = _sleeper + self._slept_for: float = 0.0 + + def _push_next_tick(self) -> None: + self._next_tick = datetime.now() + self._tick_interval + + @staticmethod + def _seconds_until(t: datetime) -> float: + td: timedelta = t - datetime.now() + return max(td.total_seconds(), 0.0) + + def _sleep(self, seconds: float) -> None: + self._sleeper(seconds) + self._slept_for += seconds + + def tick(self) -> None: + """ + You should call this method before making a request. + The first time will be instantaneous. + """ + self._sleep(self._seconds_until(self._next_tick)) + self._push_next_tick() + + @classmethod + def _unlimited(cls) -> "_Limiter": + return cls(tick_interval=timedelta(seconds=0)) diff --git a/flareio/api_client.py b/flareio/api_client.py index 947b6b0..a843c0e 100644 --- a/flareio/api_client.py +++ b/flareio/api_client.py @@ -11,6 +11,7 @@ import typing as t +from flareio._ratelimit import _Limiter from flareio.exceptions import TokenError from flareio.version import __version__ as _flareio_version @@ -263,3 +264,57 @@ def scroll( params["from"] = next_page if json and from_in_json: json["from"] = next_page + + def scroll_events( + self, + *, + method: t.Literal[ + "GET", + "POST", + ], + pages_url: str, + events_url: str, + params: t.Optional[t.Dict[str, t.Any]] = None, + json: t.Optional[t.Dict[str, t.Any]] = None, + _pages_limiter: t.Optional[_Limiter] = None, + _events_limiter: t.Optional[_Limiter] = None, + ) -> t.Iterator[ + t.Tuple[ + dict, + t.Optional[str], + ], + ]: + pages_limiter: _Limiter = _pages_limiter or _Limiter( + tick_interval=timedelta(seconds=1), + ) + events_limiter: _Limiter = _events_limiter or _Limiter( + tick_interval=timedelta(seconds=0.25), + ) + + pages_limiter.tick() + for page_resp in self.scroll( + method=method, + url=pages_url, + params=params, + json=json, + ): + page_resp.raise_for_status() + page_items: t.List[dict] = page_resp.json()["items"] + page_next: t.Optional[str] = page_resp.json()["next"] + + for page_item in page_items: + event_uid: str = page_item["metadata"]["uid"] + + events_limiter.tick() + event_resp: requests.Response = self.get( + url=events_url, + params={ + "uid": event_uid, + }, + ) + event_resp.raise_for_status() + event: dict = event_resp.json() + + yield event, page_next + + pages_limiter.tick() diff --git a/tests/test_api_client_scroll_events.py b/tests/test_api_client_scroll_events.py new file mode 100644 index 0000000..dfff24d --- /dev/null +++ b/tests/test_api_client_scroll_events.py @@ -0,0 +1,65 @@ +import pytest +import requests_mock + +from .utils import get_test_client + +from flareio._ratelimit import _Limiter + + +def test_scroll_events() -> None: + api_client = get_test_client() + + no_limit: _Limiter = _Limiter._unlimited() + + # This should make no http call. + with requests_mock.Mocker() as mocker: + events_iterator = api_client.scroll_events( + method="GET", + pages_url="https://api.flare.io/pages", + events_url="https://api.flare.io/events", + params={ + "from": None, + }, + _pages_limiter=no_limit, + _events_limiter=no_limit, + ) + assert len(mocker.request_history) == 0 + + # First page + with requests_mock.Mocker() as mocker: + mocker.register_uri( + "GET", + "https://api.flare.io/pages", + json={ + "items": [ + {"metadata": {"uid": "first_event_uid"}}, + ], + "next": "second_page", + }, + status_code=200, + ) + mocker.register_uri( + "GET", + "https://api.flare.io/events", + json={"event": "hello"}, + status_code=200, + ) + + item, cursor = next(events_iterator) + assert len(mocker.request_history) == 2 + assert item == {"event": "hello"} + assert cursor == "second_page" + + # Last page + with requests_mock.Mocker() as mocker: + mocker.register_uri( + "GET", + "https://api.flare.io/pages", + json={ + "items": [], + "next": None, + }, + ) + with pytest.raises(StopIteration): + next(events_iterator) + assert len(mocker.request_history) == 1 diff --git a/tests/test_ratelimit.py b/tests/test_ratelimit.py new file mode 100644 index 0000000..a0a7514 --- /dev/null +++ b/tests/test_ratelimit.py @@ -0,0 +1,46 @@ +from datetime import datetime +from datetime import timedelta + +from flareio._ratelimit import _Limiter + + +def test_limiter() -> None: + # Setup limiter + limiter: _Limiter = _Limiter( + tick_interval=timedelta(seconds=35), + _sleeper=lambda _: None, + ) + + # The first tick is instantaneous. This is so that the limiter + # never sleeps if the requests are taking longer than the interval. + t_1: datetime = limiter._next_tick + limiter.tick() + t_2: datetime = limiter._next_tick + assert limiter._slept_for == 0 + assert t_2 > t_1 + + # Second tick is delayed. + limiter.tick() + t_3: datetime = limiter._next_tick + assert t_3 > t_2 + assert limiter._slept_for > 30 # Intentionally not exact to avoid test races. + + +def test_seconds_until() -> None: + future: datetime = datetime.now() + timedelta(seconds=10) + assert _Limiter._seconds_until(future) > 5 + + +def test_seconds_until_negative() -> None: + now: datetime = datetime.now() + past = now - timedelta(seconds=10) + assert _Limiter._seconds_until(past) == 0.0 + + +def test_limiter_unlimited() -> None: + limiter: _Limiter = _Limiter._unlimited() + assert limiter._slept_for == 0 + limiter.tick() + limiter.tick() + limiter.tick() + assert limiter._slept_for == 0