From abc05ba85e9bd19307d6c292fc3193715e460e1e Mon Sep 17 00:00:00 2001
From: Kay Faraday <iexoo7ahth8s@freak.university>
Date: Sat, 2 Jul 2022 23:35:16 +0000
Subject: [PATCH] handle rate limits

---
 mirror_bot.py |  5 +++--
 pleroma.py    |  5 +++--
 utils.py      | 45 +++++++++++++++++++++++++++++++--------------
 3 files changed, 37 insertions(+), 18 deletions(-)

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 [