Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- from tenacity import retry, wait_fixed, stop_after_attempt, before_log
- # ....
- LAST_REQUESTS = {}
- LAST_REQUESTS_KEY = 0
- def httpreq(throttle_wait, **kwargs):
- throttle_wait = throttle_wait
- retry_kwargs = kwargs
- def real_decorator(function):
- if len(retry_kwargs):
- @retry(**retry_kwargs)
- async def wrapper(*args, **kwargs):
- if isinstance(LAST_REQUESTS_KEY, int):
- key = args[LAST_REQUESTS_KEY+1]
- else:
- key = kwargs[LAST_REQUESTS_KEY]
- kwargs.pop(LAST_REQUESTS_KEY)
- date = LAST_REQUESTS.get(key, None)
- if not date:
- LAST_REQUESTS[key] = datetime.datetime.now()
- else:
- if (datetime.datetime.now() - date).total_seconds() < throttle_wait:
- await asyncio.sleep(throttle_wait)
- else:
- LAST_REQUESTS[key] = datetime.datetime.now()
- return await function(*args, **kwargs)
- return wrapper
- else:
- async def wrapper(*args, **kwargs):
- if isinstance(LAST_REQUESTS_KEY, int):
- key = args[LAST_REQUESTS_KEY + 1]
- else:
- key = kwargs[LAST_REQUESTS_KEY]
- kwargs.pop(LAST_REQUESTS_KEY)
- date = LAST_REQUESTS.get(key, None)
- if not date:
- LAST_REQUESTS[key] = datetime.datetime.now()
- else:
- if (datetime.datetime.now() - date).total_seconds() < throttle_wait:
- await asyncio.sleep(throttle_wait)
- else:
- LAST_REQUESTS[key] = datetime.datetime.now()
- return await function(*args, **kwargs)
- return wrapper
- return real_decorator
- @httpreq(throttle_wait=settings.HTTP_REQUEST.get('THROTTLE_WAIT', None),
- wait=wait_fixed(settings.HTTP_REQUEST.get('RETRY_WAIT_FIXED', None)),
- stop=stop_after_attempt(settings.HTTP_REQUEST.get('RETRY_STOP_AFTER_ATTEMPT', None)))
- async def request(self, *args, **kwargs):
- response = None
- # ...
- return response
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement