From ea2a1d9eccd8ff355ea9d141757613000b7800fe Mon Sep 17 00:00:00 2001
From: Roman Tomjak <r00m@users.noreply.github.com>
Date: Wed, 1 Nov 2017 01:08:21 +0000
Subject: [PATCH] add CBV support (#2)

* add CBV support

* add cbv test

* add readme
---
 .gitignore                    |  5 ++++-
 README.md                     | 10 ++++++++++
 redis_ratelimit/decorators.py |  7 ++++++-
 setup.py                      |  8 ++++----
 tests/tests.py                | 19 +++++++++++++++++++
 5 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/.gitignore b/.gitignore
index fc3ef70..069b2c0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,4 +2,7 @@
 .tox/
 drr/
 *.egg-info
-__pycache__
\ No newline at end of file
+__pycache__
+build/
+dist/
+README
diff --git a/README.md b/README.md
index 8b8bf89..f8ef0de 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,12 @@ A fixed window rate limiting based on Redis
 
 ---
 
+## Requirements
+
+- Python >= 3.6
+- Django >= 1.11
+- Redis
+
 ## Installation
 
 To install django-redis-ratelimit, simply:
@@ -53,3 +59,7 @@ For this example we will assume that each key takes up roughly 250 bytes and eac
 ## Notes
 
 - [Redis Rate Limiting Pattern #2](https://redis.io/commands/INCR#pattern-rate-limiter-2)
+
+## License
+
+MIT
\ No newline at end of file
diff --git a/redis_ratelimit/decorators.py b/redis_ratelimit/decorators.py
index 79eda7a..98c3ddc 100644
--- a/redis_ratelimit/decorators.py
+++ b/redis_ratelimit/decorators.py
@@ -2,6 +2,7 @@
 
 import redis
 from django.conf import settings
+from django.http import HttpRequest
 
 from redis_ratelimit.exceptions import RateLimited
 from redis_ratelimit.utils import parse_rate, build_redis_key
@@ -34,7 +35,11 @@ def ratelimit(rate=None):
     def decorator(f):
         @wraps(f)
         def decorated_function(*args, **kwargs):
-            request = args[0]
+            # CBV support
+            if isinstance(args[0], HttpRequest):
+                request = args[0]
+            else:
+                request = args[1]
             if is_rate_limited(request, rate=rate):
                 raise RateLimited("Too Many Requests")
             return f(*args, **kwargs)
diff --git a/setup.py b/setup.py
index e14ae41..cef2320 100644
--- a/setup.py
+++ b/setup.py
@@ -7,13 +7,13 @@
 here = path.abspath(path.dirname(__file__))
 
 # Get the long description from the README file
-with open(path.join(here, 'README.md'), encoding='utf-8') as f:
+with open(path.join(here, 'README'), encoding='utf-8') as f:
     long_description = f.read()
 
 setup(
     name='django-redis-ratelimit',
-    version='0.0.1',
-    description='A sliding window rate limiting based on Redis',
+    version='0.1.1',
+    description='A fixed window rate limiting based on Redis',
     long_description=long_description,
     url='https://github.com/r00m/django-redis-ratelimit',
     author='Roman Tomjak',
@@ -27,7 +27,7 @@
         'License :: OSI Approved :: MIT License',
         'Programming Language :: Python :: 3.6',
     ],
-    keywords='django redis rate-limit',
+    keywords='django redis rate-limit ratelimit',
     packages=find_packages(exclude=['contrib', 'docs', 'tests']),
     install_requires=['Django >= 1.11', 'redis'],
 )
diff --git a/tests/tests.py b/tests/tests.py
index 18d4959..c672185 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -1,6 +1,7 @@
 from django.conf.urls import url
 from django.test import RequestFactory, TestCase
 from django.test.utils import override_settings
+from django.views import View
 
 from redis_ratelimit import ratelimit
 from redis_ratelimit.exceptions import RateLimited
@@ -51,3 +52,21 @@ class DynamicUrlPattern:
             with self.assertRaises(RateLimited):
                 req = factory.get('/')
                 view(req)
+
+    def test_cbv_decorator(self):
+        class Cbv(View):
+            @ratelimit(rate='5/s')
+            def get(self, request):
+                return True
+
+        class DynamicUrlPattern:
+            urlpatterns = [url(r'', Cbv.as_view())]
+
+        with override_settings(ROOT_URLCONF=DynamicUrlPattern):
+            for _ in range(5):
+                req = factory.get('/')
+                Cbv.as_view()(req)
+
+            with self.assertRaises(RateLimited):
+                req = factory.get('/')
+                Cbv.as_view()(req)