Merge pull request #12 from Jude188/feature/decorator

Implement function decorator

doc/source/decorators.rst

@@ -0,0 +1,46 @@
+==========================
+ Rush's Throttle Decorator
+==========================
+
+:class:`~rush.decorator.ThrottleDecorator` is an inferace which allows
+Rush's users to limit calls to a function using a decorator. Both
+synchronous and asynchronous functions are supported.
+
+
+.. autoclass:: rush.decorator.ThrottleDecorator
+   :members:
+
+
+Example
+=======
+
+.. code-block:: python
+
+    from rush import decorator
+    from rush import exceptions
+    from rush import quota
+    from rush import throttle
+    from rush.limiters import periodic
+    from rush.stores import dictionary
+
+
+    t = throttle.Throttle(
+    limiter=periodic.PeriodicLimiter(
+        store=dictionary.DictionaryStore()
+        ),
+        rate=quota.Quota.per_second(
+            count=1,
+        ),
+    )
+
+    @decorator.ThrottleDecorator(throttle=t)
+    def ratelimited_func():
+        return True
+    try:
+        for _ in range(2):
+            ratelimited_func()
+    except exceptions.ThrottleExceeded as e:
+        limit_result = e.result
+        print(limit_result.limited)  # => True
+        print(limit_result.remaining)  # => 0
+        print(limit_result.reset_after)  # => ~0:00:01

doc/source/index.rst

@@ -94,6 +94,7 @@ Table of Contents
    :maxdepth: 2
 
    throttles
+   decorators
    limiters
    storage
    examples

src/rush/decorator.py

@@ -0,0 +1,146 @@
+"""Throttle decorator public interface."""
+import asyncio
+import functools
+import inspect
+import time
+import typing
+
+import attr
+
+from rush import exceptions as rexc
+from rush import result as res
+from rush import throttle as thr
+
+
+@attr.s
+class ThrottleDecorator:
+    """The class that acts as a decorator used to throttle function calls.
+
+    This class requires an intantiated throttle with which to limit function
+    invocations.
+
+    .. attribute:: throttle
+
+        The :class:`~rush.throttle.Throttle` which should be used to limit
+        decorated functions.
+    """
+
+    throttle: thr.Throttle = attr.ib()
+
+    def _check(self, key: str, thr: thr.Throttle) -> res.RateLimitResult:
+        result = self.throttle.check(key=key, quantity=1)
+        if result.limited:
+            raise rexc.ThrottleExceeded("Rate-limit exceeded", result=result)
+        return result
+
+    def __call__(self, func: typing.Callable) -> typing.Callable:
+        """Wrap a function with a Throttle.
+
+        :param callable func:
+            The function to decorate.
+        :return:
+            Decorated function.
+        :rtype:
+            :class:`~typing.Callable`
+        """
+        key = func.__name__
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs) -> typing.Callable:
+                """Throttle the decorated function.
+
+                Extend the behaviour of the decorated function, forwarding
+                function calls if the throttle allows. The decorator will
+                raise an exception if the function cannot be called so the
+                caller may implement a retry strategy.
+
+                :param args:
+                    non-keyword arguments to pass to the decorated function.
+                :param kwargs:
+                    keyworded arguments to pass to the decorated function.
+                :raises:
+                    `~rush.exceptions.ThrottleExceeded`
+                """
+                self._check(key=key, thr=self.throttle)
+                return await func(*args, **kwargs)
+
+        else:
+
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs) -> typing.Callable:
+                """Throttle the decorated function.
+
+                Extend the behaviour of the decorated function, forwarding
+                function calls if the throttle allows. The decorator will
+                raise an exception if the function cannot be called so the
+                caller may implement a retry strategy.
+
+                :param args:
+                    non-keyword arguments to pass to the decorated function.
+                :param kwargs:
+                    keyworded arguments to pass to the decorated function.
+                :raises:
+                    `~rush.exceptions.ThrottleExceeded`
+                """
+                self._check(key=key, thr=self.throttle)
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    def sleep_and_retry(self, func: typing.Callable) -> typing.Callable:
+        """Wrap function with a naive sleep and retry strategy.
+
+        :param Callable func:
+            The :class:`~typing.Callable` to decorate.
+        :return:
+            Decorated function.
+        :rtype:
+            :class:`~typing.Callable`
+        """
+        throttled_func = self(func)
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs) -> typing.Callable:
+                """Perform naive sleep and retry strategy.
+
+                Call the throttled function. If the function raises a
+                `ThrottleExceeded` exception sleep for the recommended
+                time and retry.
+
+                :param args:
+                    non-keyword arguments to pass to the decorated function.
+                :param kwargs:
+                    keyworded arguments to pass to the decorated function.
+                """
+                while True:
+                    try:
+                        return await throttled_func(*args, **kwargs)
+                    except rexc.ThrottleExceeded as e:
+                        await asyncio.sleep(
+                            e.result.retry_after.total_seconds()
+                        )
+
+        else:
+
+            @functools.wraps(throttled_func)
+            def wrapper(*args, **kwargs) -> typing.Callable:
+                """Perform naive sleep and retry strategy.
+
+                Call the throttled function. If the function raises a
+                `ThrottleExceeded` exception sleep for the recommended
+                time and retry.
+
+                :param args:
+                    non-keyword arguments to pass to the decorated function.
+                :param kwargs:
+                    keyworded arguments to pass to the decorated function.
+                """
+                while True:
+                    try:
+                        return throttled_func(*args, **kwargs)
+                    except rexc.ThrottleExceeded as e:
+                        time.sleep(e.result.retry_after.total_seconds())
+
+        return wrapper

src/rush/exceptions.py

@@ -1,4 +1,5 @@
 """Exceptions for the rush library."""
+from rush import result
 
 
 class RushError(Exception):
@@ -40,3 +41,12 @@ def __init__(self, message, *, url, error):
         super().__init__(message)
         self.url = url
         self.error = error
+
+
+class ThrottleExceeded(Exception):
+    """The rate-limit has been exceeded."""
+
+    def __init__(self, message, *, result: result.RateLimitResult) -> None:
+        """Handle extra arguments for easier access by users."""
+        super().__init__(message)
+        self.result = result

test/unit/test_throttle_decorator.py

@@ -0,0 +1,114 @@
+"""Tests for our decorator module."""
+import asyncio
+import datetime
+
+import mock
+import pytest
+
+from rush import decorator
+from rush import exceptions
+
+
+class TestThrottleDecorator:
+    """Tests for our ThrottleDecorator class."""
+
+    def test_call_sync_non_limited(self):
+        """Verify that we can call a synchronous function."""
+        res = mock.Mock()
+        res.limited = False
+        t = mock.Mock()
+        t.check.return_value = res
+
+        @decorator.ThrottleDecorator(throttle=t)
+        def test_func():
+            return True
+
+        assert test_func() is True
+
+    def test_call_sync_limited(self):
+        """Verify that a synchronous function is throttled."""
+        res = mock.Mock()
+        res.limited = True
+        t = mock.Mock()
+        t.check.return_value = res
+
+        @decorator.ThrottleDecorator(throttle=t)
+        def test_func():
+            return True
+
+        with pytest.raises(exceptions.ThrottleExceeded):
+            test_func()
+
+    def test_call_async_non_limited(self):
+        """Verify that we can call an asynchronous function."""
+        res = mock.Mock()
+        res.limited = False
+        t = mock.Mock()
+        t.check.return_value = res
+
+        @decorator.ThrottleDecorator(throttle=t)
+        async def test_func():
+            return True
+
+        loop = asyncio.get_event_loop()
+        assert loop.run_until_complete(test_func()) is True
+
+    def test_call_async_limited(self):
+        """Verify that an asynchronous function is throttled."""
+        res = mock.Mock()
+        res.limited = True
+        t = mock.Mock()
+        t.check.return_value = res
+
+        @decorator.ThrottleDecorator(throttle=t)
+        async def test_func():
+            return True
+
+        loop = asyncio.get_event_loop()
+        with pytest.raises(exceptions.ThrottleExceeded):
+            loop.run_until_complete(test_func())
+
+    def test_sleep_and_retry_sync(self):
+        """Verify that a synchronous function is retried."""
+        retry_after = datetime.timedelta(seconds=0.5)
+
+        first_res = mock.Mock()
+        first_res.limited = True
+        first_res.retry_after = retry_after
+        second_res = mock.Mock()
+        second_res.limited = False
+        t = mock.Mock()
+        t.check.side_effect = [first_res, second_res]
+
+        throttle_decorator = decorator.ThrottleDecorator(throttle=t)
+
+        @throttle_decorator.sleep_and_retry
+        def test_func():
+            return datetime.datetime.now()
+
+        now = datetime.datetime.now()
+        res = test_func()
+        assert res - now > retry_after
+
+    def test_sleep_and_retry_async(self):
+        """Verify that an asynchronous function is retried."""
+        retry_after = datetime.timedelta(seconds=0.5)
+
+        first_res = mock.Mock()
+        first_res.limited = True
+        first_res.retry_after = retry_after
+        second_res = mock.Mock()
+        second_res.limited = False
+        t = mock.Mock()
+        t.check.side_effect = [first_res, second_res]
+
+        throttle_decorator = decorator.ThrottleDecorator(throttle=t)
+
+        @throttle_decorator.sleep_and_retry
+        async def test_func():
+            return datetime.datetime.now()
+
+        loop = asyncio.get_event_loop()
+        now = datetime.datetime.now()
+        res = loop.run_until_complete(test_func())
+        assert res - now > retry_after