diff --git a/.travis.yml b/.travis.yml index 4dbf250..0a166d6 100644 --- a/.travis.yml +++ b/.travis.yml @@ -1,15 +1,12 @@ --- language: python python: - - "2.7" - - "3.2" + - "3.5" sudo: false matrix: fast_finish: true - allow_failures: - - python: "3.2" script: - py.test diff --git a/README.rst b/README.rst index 1163ff0..b5422d3 100644 --- a/README.rst +++ b/README.rst @@ -69,7 +69,24 @@ To use this package simply decorate any function that makes an API call: raise Exception('API response: {}'.format(response.status_code)) return response -This function will not be able to make more then 15 API call within a 15 minute +Similarly, it is possible to decorate async functions: +.. code:: python + import aiohttp + from ratelimit import limits + + FIFTEEN_MINUTES = 900 + + @limits(calls=15, period=FIFTEEN_MINUTES) + async def curl(url): + async with aiohttp.ClientSession() as session: + async with session.request('GET', url) as response: + print(repr(response)) + chunk = await response.content.read() + print('Downloaded: %s' % len(chunk)) + + + +These functions will not be able to make more then 15 API call within a 15 minute time period. The arguments passed into the decorator describe the number of function @@ -77,7 +94,7 @@ invocation allowed over a specified time period (in seconds). If no time period is specified then it defaults to 15 minutes (the time window imposed by Twitter). -If a decorated function is called more times than that allowed within the +If a decorated function is called (or awaited in async case) more times than that allowed within the specified time period then a ``ratelimit.RateLimitException`` is raised. This may be used to implement a retry strategy such as an `expoential backoff `_ @@ -103,7 +120,8 @@ may be used to implement a retry strategy such as an `expoential backoff Alternatively to cause the current thread to sleep until the specified time period has ellapsed and then retry the function use the ``sleep_and_retry`` decorator. This ensures that every function invocation is successful at the -cost of halting the thread. +cost of halting the thread. This decorator also works with async functions, +causing it to `asyncio.sleep` until cooldown time expires. .. code:: python diff --git a/ratelimit/decorators.py b/ratelimit/decorators.py index b290500..0625d88 100644 --- a/ratelimit/decorators.py +++ b/ratelimit/decorators.py @@ -11,10 +11,13 @@ import time import sys import threading +import inspect +import asyncio from ratelimit.exception import RateLimitException from ratelimit.utils import now + class RateLimitDecorator(object): ''' Rate limit decorator class. @@ -42,6 +45,26 @@ def __init__(self, calls=15, period=900, clock=now(), raise_on_limit=True): # Add thread safety. self.lock = threading.RLock() + def _check_limit_exceeded(self): + with self.lock: + period_remaining = self.__period_remaining() + + # If the time window has elapsed then reset. + if period_remaining <= 0: + self.num_calls = 0 + self.last_reset = self.clock() + + # Increase the number of attempts to call the function. + self.num_calls += 1 + + # If the number of attempts to call the function exceeds the + # maximum then raise an exception. + if self.num_calls > self.clamped_calls: + if self.raise_on_limit: + raise RateLimitException('too many calls', period_remaining) + return True + return False + def __call__(self, func): ''' Return a wrapped function that prevents further function invocations if @@ -64,26 +87,23 @@ def wrapper(*args, **kargs): :param kargs: keyworded variable length argument list to the decorated function. :raises: RateLimitException ''' - with self.lock: - period_remaining = self.__period_remaining() - - # If the time window has elapsed then reset. - if period_remaining <= 0: - self.num_calls = 0 - self.last_reset = self.clock() + if self._check_limit_exceeded(): + return + return func(*args, **kargs) - # Increase the number of attempts to call the function. - self.num_calls += 1 + @wraps(func) + async def async_wrapper(*args, **kargs): + ''' + Does the same thing as `wrapper` but in async manner + ''' + if self._check_limit_exceeded(): + return + return await func(*args, **kargs) - # If the number of attempts to call the function exceeds the - # maximum then raise an exception. - if self.num_calls > self.clamped_calls: - if self.raise_on_limit: - raise RateLimitException('too many calls', period_remaining) - return - - return func(*args, **kargs) - return wrapper + if inspect.iscoroutinefunction(func): + return async_wrapper + else: + return wrapper def __period_remaining(self): ''' @@ -95,6 +115,7 @@ def __period_remaining(self): elapsed = self.clock() - self.last_reset return self.period - elapsed + def sleep_and_retry(func): ''' Return a wrapped function that rescues rate limit exceptions, sleeping the @@ -118,4 +139,19 @@ def wrapper(*args, **kargs): return func(*args, **kargs) except RateLimitException as exception: time.sleep(exception.period_remaining) - return wrapper + + @wraps(func) + async def async_wrapper(*args, **kargs): + ''' + Does the same thing as `wrapper` but in async manner + ''' + while True: + try: + return await func(*args, **kargs) + except RateLimitException as exception: + await asyncio.sleep(exception.period_remaining) + + if inspect.iscoroutinefunction(func): + return async_wrapper + else: + return wrapper diff --git a/requirements.txt b/requirements.txt index 911a822..2b1a95e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,5 @@ -pytest==2.6.4 +pytest==5.3.2 pytest-cov==2.5.1 -pylint==1.7.2 +pytest-asyncio==0.10.0 +pylint==2.3.1 +aiounittest==1.2.1 diff --git a/setup.py b/setup.py index 4052f9c..d01f75e 100644 --- a/setup.py +++ b/setup.py @@ -17,6 +17,7 @@ def readme(): url='https://github.com/tomasbasham/ratelimit', license='MIT', packages=['ratelimit'], + python_requires='>= 3.5', install_requires=[], keywords=[ 'ratelimit', diff --git a/tests/unit/async_decorator_test.py b/tests/unit/async_decorator_test.py new file mode 100644 index 0000000..a49b955 --- /dev/null +++ b/tests/unit/async_decorator_test.py @@ -0,0 +1,69 @@ +''' + +''' + +import aiounittest +import pytest +import inspect + +from ratelimit import limits, RateLimitException, sleep_and_retry +from tests import clock +from unittest.mock import patch + + +async def async_func_to_test(): + ''' + Basic async function returning True + ''' + return True + + +def sync_func_to_test(): + ''' + Basic sync function returning True + ''' + return True + + +class TestDecorator(aiounittest.AsyncTestCase): + ''' + Tests for asyncio integration with ratelimit + ''' + + @pytest.mark.asyncio + async def test_takes_sync_and_async_func(self): + ''' + Checks if sync/async wrapper selection works + ''' + limited_sync = limits(calls=1, period=10, clock=clock)(sync_func_to_test) + self.assertFalse(inspect.iscoroutinefunction(limited_sync)) + self.assertTrue(limited_sync()) + + limited_async = limits(calls=1, period=10, clock=clock)(async_func_to_test) + self.assertTrue(inspect.iscoroutinefunction(limited_async)) + self.assertTrue(await limited_async()) + + @pytest.mark.asyncio + async def test_async_function_raises(self): + ''' + Checks if async limiting raises RateLimitException same to sync method + ''' + with self.assertRaises(RateLimitException): + limited_async = limits(calls=1, period=10, clock=clock)(async_func_to_test) + await limited_async() + await limited_async() + + async def _mock_sleep(self, *args, **kwargs): + clock.increment() + + @pytest.mark.asyncio + async def test_sleep_and_retry_async(self): + period = 0.1 + sleep_mock = patch('ratelimit.decorators.asyncio.sleep').start() + sleep_mock.side_effect = self._mock_sleep + fun = sleep_and_retry(limits(calls=1, period=period, clock=clock)(async_func_to_test)) + self.assertTrue(inspect.iscoroutinefunction(fun)) + + await fun() + await fun() + sleep_mock.assert_called_once_with(period)