Skip to content
This repository was archived by the owner on Oct 12, 2022. It is now read-only.

Commit d329757

Browse files
proguy914629botproguy914629bot
committed
Handle ratelimit and bump version
1 parent 5f94c63 commit d329757

File tree

5 files changed

+144
-55
lines changed

5 files changed

+144
-55
lines changed

‎openrobot/api_wrapper/__init__.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,4 @@
77

88
from . import _async, _sync, translate, results, error, speech
99

10-
__version__ = '0.3.0.1'
10+
__version__ = '0.3.0.2'

‎openrobot/api_wrapper/_async.py

+72-32
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,15 @@ class AsyncClient:
3232
The loop to be used. Defaults to :meth:`asyncio.get_event_loop`.
3333
ignore_warning: Optional[:class:`bool`]
3434
Ignores the ``I-Am-Testing`` Warning. Defaults to ``False``.
35+
handle_ratelimit: Optional[:class:`bool`]
36+
Handles the ratelimit. If this is ``False``, then it raises
37+
:exc:`TooManyRequests`. Else, it will sleep for `Retry-After`.
38+
Defaults to ``True``.
39+
tries: Optional[:class:`int`]
40+
The number of tries to execute a request to the API This is to.
41+
handle 429s. This does not affect anything if ``handle_ratelimit``
42+
is ``False``. If this is ``None``, it will go infinitely and you
43+
might get Temp-Banned by Cloudflare. Defaults to ``5``.
3544
3645
Attributes
3746
----------
@@ -43,21 +52,25 @@ class AsyncClient:
4352
The session used. ``None`` if not specified.
4453
"""
4554

46-
def __init__(self, token: str = 'I-Am-Testing', *, session: aiohttp.ClientSession = None, loop: asyncio.AbstractEventLoop = None, ignore_warning: bool = False):
55+
def __init__(self, token: str = 'I-Am-Testing', *, session: aiohttp.ClientSession = None, loop: asyncio.AbstractEventLoop = None, ignore_warning: bool = False, handle_ratelimit: bool = True, tries: int = 5):
4756
if not token:
4857
raise NoTokenProvided()
4958
elif token == 'I-Am-Testing' and not ignore_warning:
5059
warnings.warn('Using I-Am-Testing token will only let you 5 requests/day (UTC based, will reset on 00:00 UTC) for all `/api` methods and will raise an `openrobot.api_wrapper.error.Forbidden` after you have reached your limit.')
5160

52-
self.token = str(token)
61+
self.token: str = str(token)
5362

54-
self.loop = loop or asyncio.get_event_loop()
63+
self.loop: asyncio.AbstractEventLoop = loop or asyncio.get_event_loop()
5564

56-
self.session = session if isinstance(session, aiohttp.ClientSession) else None
65+
self.session: typing.Optional[aiohttp.ClientSession] = session if isinstance(session, aiohttp.ClientSession) else None
5766

5867
if self.session:
5968
self.session._loop = self.loop
6069

70+
self.handle_ratelimit: bool = handle_ratelimit
71+
72+
self.tries: int = tries
73+
6174
# Important and internal methods, but should be used un-regularly by the User itself.
6275

6376
def _get_authorization_headers(self, token: str = None, *, header = True):
@@ -90,35 +103,12 @@ async def _request(self, method: str, url: str, **kwargs) -> typing.Union[dict,
90103
url = ('https://api.openrobot.xyz/api' + url)
91104
else:
92105
raise TypeError('URL is not a valid HTTP/HTTPs URL.')
93-
94-
if self.session:
95-
async with self.session.request(method, url, **kwargs) as resp:
96-
js = await resp.json()
97-
if resp.status in return_on:
98-
if raw:
99-
return resp
100-
else:
101-
return js
102-
elif resp.status == 403:
103-
raise Forbidden(resp, js)
104-
elif resp.status == 400:
105-
raise BadRequest(resp, js)
106-
elif resp.status == 500:
107-
raise InternalServerError(resp, js)
108-
elif 200 <= resp.status < 300:
109-
if raw:
110-
return resp
111-
else:
112-
return js
113-
else:
114-
cls = OpenRobotAPIError(js)
115-
cls.raw = js
116-
cls.response = resp
117106

118-
raise cls
119-
else:
120-
async with aiohttp.ClientSession(loop=self.loop) as sess:
121-
async with sess.request(method, url, **kwargs) as resp:
107+
tries = int(self.tries) if self.tries is not None else None
108+
109+
while tries is None or tries > 0:
110+
if self.session:
111+
async with self.session.request(method, url, **kwargs) as resp:
122112
js = await resp.json()
123113
if resp.status in return_on:
124114
if raw:
@@ -131,6 +121,17 @@ async def _request(self, method: str, url: str, **kwargs) -> typing.Union[dict,
131121
raise BadRequest(resp, js)
132122
elif resp.status == 500:
133123
raise InternalServerError(resp, js)
124+
elif resp.status == 429:
125+
if not self.handle_ratelimit:
126+
raise TooManyRequests(resp, js)
127+
128+
try:
129+
await asyncio.sleep(resp.headers['Retry-After'])
130+
except KeyError as e:
131+
raise KeyError('Retry-After header is not present.') from e
132+
133+
if tries:
134+
tries -= 1
134135
elif 200 <= resp.status < 300:
135136
if raw:
136137
return resp
@@ -142,6 +143,45 @@ async def _request(self, method: str, url: str, **kwargs) -> typing.Union[dict,
142143
cls.response = resp
143144

144145
raise cls
146+
else:
147+
async with aiohttp.ClientSession(loop=self.loop) as sess:
148+
async with sess.request(method, url, **kwargs) as resp:
149+
js = await resp.json()
150+
if resp.status in return_on:
151+
if raw:
152+
return resp
153+
else:
154+
return js
155+
elif resp.status == 403:
156+
raise Forbidden(resp, js)
157+
elif resp.status == 400:
158+
raise BadRequest(resp, js)
159+
elif resp.status == 500:
160+
raise InternalServerError(resp, js)
161+
elif resp.status == 429:
162+
if not self.handle_ratelimit:
163+
raise TooManyRequests(resp, js)
164+
165+
try:
166+
await asyncio.sleep(resp.headers['Retry-After'])
167+
except KeyError as e:
168+
raise KeyError('Retry-After header is not present.') from e
169+
170+
if tries:
171+
tries -= 1
172+
elif 200 <= resp.status < 300:
173+
if raw:
174+
return resp
175+
else:
176+
return js
177+
else:
178+
cls = OpenRobotAPIError(js)
179+
cls.raw = js
180+
cls.response = resp
181+
182+
raise cls
183+
184+
raise TooManyRequests(resp, js)
145185

146186
# Methods to query to API:
147187

‎openrobot/api_wrapper/_sync.py

+59-21
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import warnings
44
import io
55
import typing
6+
import time
67
from .error import *
78
from .results import *
89
from .translate import Translate
@@ -27,20 +28,33 @@ class SyncClient:
2728
to ``I-Am-Testing``.
2829
ignore_warning: Optional[:class:`bool`]
2930
Ignores the ``I-Am-Testing`` Warning. Defaults to ``False``.
31+
handle_ratelimit: Optional[:class:`bool`]
32+
Handles the ratelimit. If this is ``False``, then it raises
33+
:exc:`TooManyRequests`. Else, it will sleep for `Retry-After`.
34+
Defaults to ``True``.
35+
tries: Optional[:class:`int`]
36+
The number of tries to execute a request to the API This is to.
37+
handle 429s. This does not affect anything if ``handle_ratelimit``
38+
is ``False``. If this is ``None``, it will go infinitely and you
39+
might get Temp-Banned by Cloudflare. Defaults to ``5``.
3040
3141
Attributes
3242
----------
3343
token: :class:`str`
3444
The token used to authorize to the API.
3545
"""
3646

37-
def __init__(self, token: str = 'I-Am-Testing', *, ignore_warning = False):
47+
def __init__(self, token: str = 'I-Am-Testing', *, ignore_warning = False, handle_ratelimit: bool = True, tries: int = 5):
3848
if not token:
3949
raise NoTokenProvided()
4050
elif token == 'I-Am-Testing' and not ignore_warning:
4151
warnings.warn('Using I-Am-Testing token will only let you 5 requests/day (UTC based, will reset on 00:00 UTC) for all `/api` methods and will raise an `openrobot.api_wrapper.error.Forbidden` after you have reached your limit.')
4252

43-
self.token = str(token)
53+
self.token: str = str(token)
54+
55+
self.handle_ratelimit: bool = handle_ratelimit
56+
57+
self.tries: int = tries
4458

4559
# Important and internal methods, but should be used un-regularly by the User itself.
4660

@@ -61,6 +75,8 @@ def _request(self, method: str, url: str, **kwargs):
6175
else:
6276
kwargs['headers'] = headers
6377

78+
return_on = kwargs.pop('return_on', [])
79+
6480
if not url.startswith('/'):
6581
url = url[1:]
6682

@@ -74,27 +90,49 @@ def _request(self, method: str, url: str, **kwargs):
7490
else:
7591
raise TypeError('URL is not a valid HTTP/HTTPs URL.')
7692

77-
with requests.Session() as s:
78-
r = s.request(method, url, **kwargs)
79-
80-
js = r.json()
81-
if r.status_code == 403:
82-
raise Forbidden(r, js)
83-
elif r.status_code == 400:
84-
raise BadRequest(r, js)
85-
elif r.status_code == 500:
86-
raise InternalServerError(r, js)
87-
elif 200 <= r.status_code < 300:
88-
if raw:
89-
return r
93+
tries = int(self.tries) if self.tries is not None else None
94+
95+
while tries is None or tries > 0:
96+
with requests.Session() as s:
97+
r = s.request(method, url, **kwargs)
98+
99+
js = r.json()
100+
101+
if r.status_code in return_on:
102+
if raw:
103+
return r
104+
else:
105+
return js
106+
if r.status_code == 403:
107+
raise Forbidden(r, js)
108+
elif r.status_code == 400:
109+
raise BadRequest(r, js)
110+
elif r.status_code == 500:
111+
raise InternalServerError(r, js)
112+
elif r.status_code == 429:
113+
if not self.handle_ratelimit:
114+
raise TooManyRequests(r, js)
115+
116+
try:
117+
time.sleep(r.headers['Retry-After'])
118+
except KeyError as e:
119+
raise KeyError('Retry-After header is not present.') from e
120+
121+
if tries:
122+
tries -= 1
123+
elif 200 <= r.status_code < 300:
124+
if raw:
125+
return r
126+
else:
127+
return js
90128
else:
91-
return js
92-
else:
93-
cls = OpenRobotAPIError(js)
94-
cls.raw = js
95-
cls.response = r
129+
cls = OpenRobotAPIError(js)
130+
cls.raw = js
131+
cls.response = r
132+
133+
raise cls
96134

97-
raise cls
135+
raise TooManyRequests(r, js)
98136

99137
# Methods to query to API:
100138

‎openrobot/api_wrapper/error.py

+11
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,17 @@ def __init__(self, resp, json_error):
3535

3636
super().__init__(self.message)
3737

38+
class TooManyRequests(OpenRobotAPIError):
39+
"""Recieved an 429 Error code from the API."""
40+
def __init__(self, resp, json_error):
41+
self.raw: dict = json_error
42+
self.response: typing.Union[aiohttp.ClientResponse, requests.Response] = resp
43+
self.message: str = json_error['message']
44+
self.error_code: int = json_error['error']['code']
45+
self.retry_after: int = resp.headers.get('Retry-After', None)
46+
47+
super().__init__(self.message)
48+
3849
class NoTokenProvided(OpenRobotAPIError):
3950
"""No token was provided."""
4051
def __init__(self):

‎pyproject.toml

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[tool.poetry]
22
name = "OpenRobot-API-Wrapper"
3-
version = "0.3.0.1"
3+
version = "0.3.0.2"
44
description = "A API Wrapper for the official OpenRobot API (https://api.openrobot.xyz)."
55
authors = ["OpenRobot Packages <[email protected]>", "proguy914629 <[email protected]>"]
66
license = "MIT"

0 commit comments

Comments
 (0)