More protection and rate limit exposure

disqusapi/__init__.py

@@ -12,6 +12,7 @@
 except:
     __version__ = 'unknown'
 
+from datetime import datetime
 import httplib
 import os.path
 import simplejson
@@ -43,7 +44,16 @@ def __str__(self):
 class InvalidAccessToken(APIError):
     pass
 
+
+class InternalServerError(APIError):
+    code = 15
+
+    def __init__(self, message):
+        self.message = message
+
+
 ERROR_MAP = {
+    15: InternalServerError,
     18: InvalidAccessToken,
 }
 
@@ -145,6 +155,7 @@ def _request(self, **kwargs):
         response = conn.getresponse()
         # Let's coerce it to Python
         data = api.formats[format](response.read())
+        api.ratelimit = RateLimit.from_response(response)
 
         if response.status != 200:
             raise ERROR_MAP.get(data['code'], APIError)(data['code'], data['response'])
@@ -154,9 +165,38 @@ def _request(self, **kwargs):
         return data['response']
 
 
+class RateLimit(object):
+
+    def __init__(self, limit=0, remaining=0, reset='now'):
+        if reset == 'now':
+            reset = datetime.utcnow()
+        else:
+            reset = datetime.fromtimestamp(float(reset))
+        self.limit = limit
+        self.remaining = remaining
+        self.reset = reset
+
+    @classmethod
+    def from_response(cls, response):
+        limit = response.getheader('X-Ratelimit-Limit')
+        limit = int(limit) if limit else 0
+        remaining = response.getheader('X-Ratelimit-Remaining')
+        remaining = int(remaining) if remaining else 0
+        reset = response.getheader('X-Ratelimit-Reset', 'now')
+        return cls(limit, remaining, reset)
+
+
+def format_json(json):
+    try:
+        result = simplejson.loads(json)
+    except simplejson.JSONDecodeError:
+        raise InternalServerError('Expected json, received: ' % json)
+    return result
+
+
 class DisqusAPI(Resource):
     formats = {
-        'json': lambda x: simplejson.loads(x),
+        'json': format_json,
     }
 
     def __init__(self, secret_key=None, public_key=None, format='json', version='3.0', **kwargs):
@@ -166,6 +206,7 @@ def __init__(self, secret_key=None, public_key=None, format='json', version='3.0
             warnings.warn('You should pass ``public_key`` in addition to your secret key.')
         self.format = format
         self.version = version
+        self.ratelimit = RateLimit()
         super(DisqusAPI, self).__init__(self)
 
     def _request(self, **kwargs):

disqusapi/tests.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 import mock
 import os
 import unittest
@@ -13,13 +14,19 @@ def wrapped(func):
     return wrapped
 
 class MockResponse(object):
-    def __init__(self, body, status=200):
+    def __init__(self, body, status=200, headers=None):
         self.body = body
         self.status = status
+        if headers is None:
+            headers = {}
+        self.headers = headers
 
     def read(self):
         return self.body
 
+    def getheader(self, key, default=None):
+        return self.headers.get(key, default)
+
 class DisqusAPITest(unittest.TestCase):
     API_SECRET = 'b'*64
     API_PUBLIC = 'c'*64
@@ -81,5 +88,27 @@ def iter_results():
                     iterator.next()
         self.assertEquals(n, 99)
 
+
+class RateLimitTest(unittest.TestCase):
+
+    def test_defaults(self):
+        rl = disqusapi.RateLimit()
+        assert rl.remaining == 0
+        assert rl.limit == 0
+        assert rl.reset <= datetime.utcnow()
+
+    def test_from_response(self):
+        timestamp = '1399359600'
+        resp = MockResponse('Hello World', headers={
+            'X-Ratelimit-Remaining': '123',
+            'X-Ratelimit-Limit': '1000',
+            'X-Ratelimit-Reset': timestamp,
+        })
+        rl = disqusapi.RateLimit.from_response(resp)
+        assert rl.remaining == 123
+        assert rl.limit == 1000
+        assert rl.reset == datetime.fromtimestamp(float(timestamp))
+
+
 if __name__ == '__main__':
     unittest.main()