Feature/asyncio decorators

.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

README.rst

@@ -69,15 +69,32 @@ 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
 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
 <https://pypi.org/project/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
 

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

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

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',

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)