Merge pull request #253 from bbolli/limit-retry

[z_archive/twitter.git] / twitter / api.py

diff --git a/twitter/api.py b/twitter/api.py
index fd33edd343e3ec6dc84ce17d45b7bacfff97f14c..059bb62f9f35eaf46eda7c860b5b127bb67921c2 100644 (file)
--- a/twitter/api.py
+++ b/twitter/api.py
@@ -1,5 +1,5 @@
 # encoding: utf-8
-from __future__ import unicode_literals
+from __future__ import unicode_literals, print_function
 
 try:
     import urllib.request as urllib_request
@@ -19,6 +19,7 @@ from .auth import NoAuth
 import re
 import sys
 import gzip
+from time import sleep, time
 
 try:
     import http.client as http_client
@@ -139,9 +140,11 @@ def method_for_uri(uri):
 
 class TwitterCall(object):
 
+    TWITTER_UNAVAILABLE_WAIT = 30  # delay after HTTP codes 502, 503 or 504
+
     def __init__(
             self, auth, format, domain, callable_cls, uri="",
-            uriparts=None, secure=True, timeout=None, gzip=False):
+            uriparts=None, secure=True, timeout=None, gzip=False, retry=False):
         self.auth = auth
         self.format = format
         self.domain = domain
@@ -151,6 +154,7 @@ class TwitterCall(object):
         self.secure = secure
         self.timeout = timeout
         self.gzip = gzip
+        self.retry = retry
 
     def __getattr__(self, k):
         try:
@@ -160,7 +164,7 @@ class TwitterCall(object):
                 return self.callable_cls(
                     auth=self.auth, format=self.format, domain=self.domain,
                     callable_cls=self.callable_cls, timeout=self.timeout,
-                    secure=self.secure, gzip=self.gzip,
+                    secure=self.secure, gzip=self.gzip, retry=self.retry,
                     uriparts=self.uriparts + (arg,))
             if k == "_":
                 return extend_call
@@ -273,7 +277,10 @@ class TwitterCall(object):
                     headers[k.encode('utf-8')] = headers.pop(k)
 
         req = urllib_request.Request(uriBase, body, headers)
-        return self._handle_response(req, uri, arg_data, _timeout)
+        if self.retry:
+            return self._handle_response_with_retry(req, uri, arg_data, _timeout)
+        else:
+            return self._handle_response(req, uri, arg_data, _timeout)
 
     def _handle_response(self, req, uri, arg_data, _timeout=None):
         kwargs = {}
@@ -308,6 +315,28 @@ class TwitterCall(object):
             else:
                 raise TwitterHTTPError(e, uri, self.format, arg_data)
 
+    def _handle_response_with_retry(self, req, uri, arg_data, _timeout=None):
+        retry = self.retry
+        while retry:
+            try:
+                return self._handle_response(req, uri, arg_data, _timeout)
+            except TwitterHTTPError as e:
+                if e.e.code == 429:
+                    # API rate limit reached
+                    reset = int(e.e.headers.get('X-Rate-Limit-Reset', time() + 30))
+                    delay = int(reset - time() + 2)  # add some extra margin
+                    print("API rate limit reached; waiting for %ds..." % delay, file=sys.stderr)
+                elif e.e.code in (502, 503, 504):
+                    delay = self.TWITTER_UNAVAILABLE_WAIT
+                    print("Service unavailable; waiting for %ds..." % delay, file=sys.stderr)
+                else:
+                    raise
+                if isinstance(retry, int):
+                    if retry <= 0:
+                        raise
+                    retry -= 1
+                sleep(delay)
+
 
 class Twitter(TwitterCall):
     """
@@ -417,7 +446,7 @@ class Twitter(TwitterCall):
     def __init__(
             self, format="json",
             domain="api.twitter.com", secure=True, auth=None,
-            api_version=_DEFAULT):
+            api_version=_DEFAULT, retry=False):
         """
         Create a new twitter API connector.
 
@@ -437,6 +466,11 @@ class Twitter(TwitterCall):
 
         `api_version` is used to set the base uri. By default it's
         '1.1'.
+
+        If `retry` is True, API rate limits will automatically be
+        handled by waiting until the next reset, as indicated by
+        the X-Rate-Limit-Reset HTTP header. If retry is an integer,
+        it defines the number of retries attempted.
         """
         if not auth:
             auth = NoAuth()
@@ -454,7 +488,7 @@ class Twitter(TwitterCall):
         TwitterCall.__init__(
             self, auth=auth, format=format, domain=domain,
             callable_cls=TwitterCall,
-            secure=secure, uriparts=uriparts)
+            secure=secure, uriparts=uriparts, retry=retry)
 
 
 __all__ = ["Twitter", "TwitterError", "TwitterHTTPError", "TwitterResponse"]