Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions flareio/_ratelimit.py
Original file line number Diff line number Diff line change
@@ -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))
55 changes: 55 additions & 0 deletions flareio/api_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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()
65 changes: 65 additions & 0 deletions tests/test_api_client_scroll_events.py
Original file line number Diff line number Diff line change
@@ -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
46 changes: 46 additions & 0 deletions tests/test_ratelimit.py
Original file line number Diff line number Diff line change
@@ -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