generator to get all of entities with start, stop. should have query support in future

Author: Rasmus Lager

ganapi/api.py

@@ -1,6 +1,8 @@
 import json
 import requests
 from gan_exception import GanException
+
+
 class Api(object):
     """
         Handles connection to the API
@@ -12,6 +14,11 @@ class Api(object):
     """
     default_base_uri = 'https://api.getanewsletter.com/v3/'
 
+    """
+        The amount of entities to get in each request while using all
+        :var batch_size int
+    """
+    batch_size = 25
     """
         Initializes the API connection
         :param string token The security token.

ganapi/entity_manager.py

@@ -1,6 +1,8 @@
 import urllib
 from helpers import PaginatedResultSet
 from api import GanException
+import urlparse
+import math
 
 class EntityManager(object):
     """
@@ -185,6 +187,61 @@ def query(self, filters=dict(), as_json=False):
             return response.json()
         return PaginatedResultSet(self, response.json())
 
+    def all(self, start=0, stop=float('inf')):
+        """
+        Method to get a generator with all or between(start, stop) entities of type.
+        Will use Api.batch_size for pagination parameter to the api.
+        When using start argument it will jump to the correct page after first request
+        if the start item is not on the first page.
+        :param start Start from entity index.
+        :param stop Stop at entity index.
+        :raises StopIteration if end of entities or too high start.
+        :raises AssertionError if start or stop are invalid.
+        :raises HTTPError if the data can not be fetched.
+        :return: generator with entities of type.
+        """
+
+        def _calc_item_page(index, total, page_size):
+            # Calculate item page based on index
+            return math.ceil(float(((index + 1) % total)) / page_size)
+
+        if start:
+            assert isinstance(start, int)
+        if stop != float('inf'):
+            assert isinstance(stop, int)
+        assert start <= stop
+
+        uri = u'{base_path}/?paginate_by={batch_size}'.format(base_path=str.rstrip(self.base_path),
+                                                              batch_size=self.api.batch_size)
+        count = None
+        count_read = 0
+        while uri and (count is None or count_read < count) and count_read <= stop:
+            results = self.api.call('GET', uri).json()
+            count = results.get('count', 0)
+
+            if start > count:
+                raise StopIteration
+
+            for entity in results.get('results', []):
+                if start <= count_read <= stop:
+                    yield self.construct_entity(entity).set_persisted()
+                count_read += 1
+
+            if self.api.batch_size < start and start > count_read + 1:
+                # Go straight to supposed page of the start item.
+                start_page = int(_calc_item_page(start, count, self.api.batch_size))
+                uri = u'{base_path}/?paginate_by={batch_size}&page={page}'.format(base_path=str.rstrip(self.base_path),
+                                                                                  batch_size=self.api.batch_size,
+                                                                                  page=start_page)
+                # Set new count_read for all skipped pages.
+                count_read = (start_page - 1) * self.api.batch_size
+            else:
+                if results.get('next'):
+                    uri = u'{base_path}/?{params}'.format(base_path=str.rstrip(self.base_path),
+                                                          params=urlparse.urlparse(results.get('next')).query)
+
+        return
+
     def delete(self, entity):
         """
          Deletes an entity.

ganapi/test_attribute_manager.py

@@ -2,14 +2,15 @@
 from api import Api
 from attribute import Attribute
 from attribute_manager import AttributeManager
-from httmock import HTTMock, all_requests
+from httmock import HTTMock, all_requests, urlmatch
 from requests import HTTPError
 
 
 class ContactManagerTest(unittest.TestCase):
     def setUp(self):
         gan_token = 'h027MapNNujPH0gV+sXAdmzZTDffHOpJEHaBtrD3NXtNqI4dT3NLXhyTwiZr7PUOGZJNSGv/b9xVyaguX0nDrONGhudPkxtl5EoXrM4SOZHswebpSy2ehh0edrGVF7dVJVZLIlRwgViY3n3/2hMQ5Njp9JFywnOy7gMeaoKw0hYLRbd+wVqvl2oOnspXwGTTcZ9Y+cdP8jIhUUoXOieXst0IXVclAHXa+K1d15gKLcpmXzK+jx14wGEmb4t8MSU'
         self.api = Api(token=gan_token)
+        self.api.batch_size = 2
         self.attribute_manager = AttributeManager(self.api)
         self.start_path = '/v3'
 
@@ -122,3 +123,44 @@ def test_get_paginated_attributes(self):
         self.assertEqual(len(paginated_result_set.entities), 6)
 
         self.assertRaises(StopIteration, paginated_result_set.next)
+
+    @all_requests
+    def get_all_attr_mock(self, url, request):
+        status_code = 200
+        if url.query == 'paginate_by=2':
+            content = '{"count":8,"next":"https://api.getanewsletter.com/v3/attributes/?page=2&paginate_by=2","previous":null,"results":[{"url":"https://api.getanewsletter.com/v3/attributes/attribute/","name":"attr0","code":"attribute","usage_count":0},{"url":"https://api.getanewsletter.com/v3/attributes/attribute2/","name":"attr1","code":"attribute2","usage_count":0}]}'
+        elif url.query in ['page=2&paginate_by=2'] :
+            content = '{"count":8,"next":"https://api.getanewsletter.com/v3/attributes/?page=3&paginate_by=2","previous":"https://api.getanewsletter.com/v3/attributes/?paginate_by=2","results":[{"url":"https://api.getanewsletter.com/v3/attributes/bu/","name":"attr2","code":"bu","usage_count":0},{"url":"https://api.getanewsletter.com/v3/attributes/bu1/","name":"attr3","code":"bu1","usage_count":0}]}'
+
+        elif url.query == 'page=3&paginate_by=2' or url.query == 'paginate_by=2&page=3':
+            content = '{"count":8,"next":"https://api.getanewsletter.com/v3/attributes/?page=4&paginate_by=2","previous":"https://api.getanewsletter.com/v3/attributes/?page=2&paginate_by=2","results":[{"url":"https://api.getanewsletter.com/v3/attributes/bu2/","name":"attr4","code":"bu2","usage_count":0},{"url":"https://api.getanewsletter.com/v3/attributes/bu3/","name":"bu3","code":"bu3","usage_count":0}]}'
+
+        elif url.query == 'page=4&paginate_by=2':
+            content = '{"count":8,"next":"https://api.getanewsletter.com/v3/attributes/?page=5&paginate_by=2","previous":"https://api.getanewsletter.com/v3/attributes/?page=3&paginate_by=2","results":[{"url":"https://api.getanewsletter.com/v3/attributes/bu4/","name":"bu4","code":"bu4","usage_count":0},{"url":"https://api.getanewsletter.com/v3/attributes/bu5/","name":"bu5","code":"bu5","usage_count":0}]}'
+        else:
+            return {'status_code': 404}
+
+        return {'content': content,
+                'status_code': status_code}
+
+    def test_get_all_attributes(self):
+        with HTTMock(self.get_all_attr_mock):
+            all_attributes = self.attribute_manager.all()
+            attrs = [attr for attr in all_attributes]
+            self.assertEqual(len(attrs), 8)
+
+    def test_get_between_attributes(self):
+        with HTTMock(self.get_all_attr_mock):
+            self.assertEqual(self.api.batch_size, 2)
+            # should get entities from page 3 and 4
+            attributes_between = self.attribute_manager.all(start=4, stop=7)
+            attrs = [attr for attr in attributes_between]
+            self.assertEqual(len(attrs), 4)
+            self.assertEqual(attrs[0].name, 'attr4')
+            self.assertTrue(isinstance(attrs[0], Attribute))
+
+    def test_get_too_high_start(self):
+        with HTTMock(self.get_all_attr_mock):
+            attributes_between = self.attribute_manager.all(start=9)
+            attrs = [attr for attr in attributes_between]
+            self.assertEqual(len(attrs), 0)