Add per-endpoint request limits

danthedaniel · Mar 28, 2017 · 29624a5 · 29624a5
1 parent cab6b73
commit 29624a5
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 7 deletions.
diff --git a/psraw/base.py b/psraw/base.py
@@ -3,7 +3,7 @@
     from urllib import urlencode
 except ImportError:
     from urllib.parse import urlencode
-from .endpoints import ENDPOINTS, BASE_ADDRESS, LIMIT_MAX, LIMIT_DEFAULT
+from .endpoints import ENDPOINTS, BASE_ADDRESS, LIMIT_DEFAULT
 
 
 def limit_chunk(limit, limit_max):
@@ -63,13 +63,15 @@ def endpoint_func(r, **kwargs):
             direction = 'after'
 
         # Loop over the API request, since multiple may be required if the
-        # specified limit is greater than LIMIT_MAX
-        for limit in limit_chunk(coerced_kwargs['limit'], LIMIT_MAX):
+        # specified limit is greater than config['limit']
+        for limit in limit_chunk(coerced_kwargs['limit'], config['limit']):
             coerced_kwargs['limit'] = limit
             query_params = urlencode(coerced_kwargs)
             url = '{}{}?{}'.format(BASE_ADDRESS, config['url'], query_params)
             data = requests.get(url).json()['data']
 
+            # import pdb; pdb.set_trace()
+
             for item in data:
                 yield config['return_type'](r, _data=item)
 

diff --git a/psraw/endpoints.py b/psraw/endpoints.py
@@ -11,7 +11,6 @@ def sort_type(value):
         raise ValueError('Value must be one of: {}'.format(directions))
 
 
-LIMIT_MAX = 500
 LIMIT_DEFAULT = 50
 BASE_ADDRESS = 'https://apiv2.pushshift.io/reddit'
 ENDPOINTS = {
@@ -24,6 +23,7 @@ def sort_type(value):
             'sort': sort_type,
             'subreddit': str
         },
+        'limit': 500,
         'return_type': praw.models.Comment,
         'url': '/comment/fetch/'
     },
@@ -36,6 +36,7 @@ def sort_type(value):
             'sort': sort_type,
             'subreddit': str
         },
+        'limit': 500,
         'return_type': praw.models.Comment,
         'url': '/search/comment/'
     },
@@ -45,6 +46,7 @@ def sort_type(value):
             'before': int,
             'limit': int
         },
+        'limit': 250,
         'return_type': praw.models.Submission,
         'url': '/submission/activity/'
     },
@@ -57,6 +59,7 @@ def sort_type(value):
             'sort': sort_type,
             'subreddit': str
         },
+        'limit': 250,
         'return_type': praw.models.Submission,
         'url': '/search/submission/'
     }

diff --git a/psraw/version.py b/psraw/version.py
@@ -1 +1 @@
-__version__ = '0.1.0'
+__version__ = '0.1.1'
diff --git a/test/test_endpoints.py b/test/test_endpoints.py
@@ -3,7 +3,7 @@
 import vcr
 import os
 import types
-from psraw.endpoints import LIMIT_MAX, ENDPOINTS
+from psraw.endpoints import ENDPOINTS
 
 
 LIMIT = 10
@@ -40,7 +40,8 @@ def test_comment_search():
 @vcr.use_cassette('test/cassettes/comment_search_large.yaml')
 def test_comment_search_2():
     """comment_search returns a list of praw.models.Comment objects"""
-    large_limit = LIMIT_MAX * 2
+    comment_search_limit = ENDPOINTS['comment_search']['limit']
+    large_limit = comment_search_limit * 2
     comments = psraw.comment_search(r, q='automoderator', limit=large_limit)
     assert isinstance(comments, types.GeneratorType)
     assert len(list(comments)) == large_limit