diff --git a/mirror_bot.py b/mirror_bot.py index 5b97c00..bfd030e 100755 --- a/mirror_bot.py +++ b/mirror_bot.py @@ -14,7 +14,7 @@ import contextlib import qtoml as toml from pleroma import Pleroma from functools import partial -from utils import suppress, loading_spinner, http_session_factory +from utils import loading_spinner, http_session_factory, HandleRateLimits JSON_CONTENT_TYPE = 'application/json' ACTIVITYPUB_CONTENT_TYPE = 'application/activity+json' @@ -37,6 +37,7 @@ class PostMirror: raise_for_status=True, ), ) + self._rl_handler = HandleRateLimits(self._http) self._ctx_stack = stack return self @@ -59,7 +60,7 @@ class PostMirror: cursor.hide() done = False while not done: - async with self._http.get(page_url) as resp: page = await resp.json() + async with self._rl_handler.request('GET', page_url) as resp: page = await resp.json() try: page_url = page['next'] except KeyError: diff --git a/pleroma.py b/pleroma.py index d07b571..b8caaf7 100644 --- a/pleroma.py +++ b/pleroma.py @@ -7,7 +7,7 @@ import hashlib import aiohttp from http import HTTPStatus from multidict import MultiDict -from utils import http_session_factory +from utils import http_session_factory, HandleRateLimits class BadRequest(Exception): pass @@ -20,6 +20,7 @@ class Pleroma: self.api_base_url = api_base_url.rstrip('/') self.access_token = access_token.strip() self._session = http_session_factory({'Authorization': 'Bearer ' + self.access_token}) + self._rl_handler = HandleRateLimits(self._session) self._logged_in_id = None async def __aenter__(self): @@ -41,7 +42,7 @@ class Pleroma: }: raise RuntimeError('stop being a chud') - async with self._session.request(method, self.api_base_url + path, **kwargs) as resp: + async with self._rl_handler.request(method, self.api_base_url + path, **kwargs) as resp: if resp.status == HTTPStatus.BAD_REQUEST: raise BadRequest((await resp.json())['error']) #resp.raise_for_status() diff --git a/utils.py b/utils.py index d1dcfaf..3f618c3 100644 --- a/utils.py +++ b/utils.py @@ -5,7 +5,7 @@ import aiohttp import platform import itertools import contextlib -from functools import wraps +from datetime import datetime, timezone def http_session_factory(headers={}, **kwargs): user_agent = ( @@ -18,21 +18,38 @@ def http_session_factory(headers={}, **kwargs): **kwargs, ) -def as_corofunc(f): - @wraps(f) - async def wrapped(*args, **kwargs): - # can't decide if i want an `anyio.sleep(0)` here. - return f(*args, **kwargs) - return wrapped +async def sleep_until(dt): + await anyio.sleep((dt - datetime.now(timezone.utc)).total_seconds()) -def as_async_cm(cls): - @wraps(cls, updated=()) # cls.__dict__ doesn't support .update() - class wrapped(cls, contextlib.AbstractAsyncContextManager): - __aenter__ = as_corofunc(cls.__enter__) - __aexit__ = as_corofunc(cls.__exit__) - return wrapped +class HandleRateLimits: + def __init__(self, http): + self.http = http -suppress = as_async_cm(contextlib.suppress) + def request(self, *args, **kwargs): + return _RateLimitContextManager(self.http, args, kwargs) + +class _RateLimitContextManager(contextlib.AbstractAsyncContextManager): + def __init__(self, http, args, kwargs): + self.http = http + self.args = args + self.kwargs = kwargs + + async def __aenter__(self): + self._request_cm = self.http.request(*self.args, **self.kwargs) + return await self._do_enter() + + async def _do_enter(self): + resp = await self._request_cm.__aenter__() + if resp.headers.get('X-RateLimit-Remaining') not in {'0', '1'}: + return resp + + print('Hit rate limit for', self.args) + await sleep_until(datetime.fromisoformat(resp.headers['X-RateLimit-Reset'])) + await self._request_cm.__aexit__(*(None,)*3) + return await self.__aenter__() + + async def __aexit__(self, *excinfo): + return await self._request_cm.__aexit__(*excinfo) def loading_spinner(): return itertools.cycle('\b' + x for x in [