Changes the storage for pagination information

fixes #232

1 files changed, 26 insertions, 8 deletions

diff --git a/mastodon/Mastodon.py b/mastodon/Mastodon.py
index 98ac72a..c22cabe 100644
--- a/mastodon/Mastodon.py
+++ b/mastodon/Mastodon.py
@@ -120,10 +120,27 @@ class AttribAccessDict(dict):
             raise AttributeError("Attribute-style access is read only")
         super(AttribAccessDict, self).__setattr__(attr, val)
 
+
 ###
-# The actual Mastodon class
+# List helper class.
+# Defined at top level so it can be pickled.
 ###
+class AttribAccessList(list):
+    def __getattr__(self, attr):
+        if attr in self:
+            return self[attr]
+        else:
+            raise AttributeError("Attribute not found: " + str(attr))
 
+    def __setattr__(self, attr, val):
+        if attr in self:
+            raise AttributeError("Attribute-style access is read only")
+        super(AttribAccessList, self).__setattr__(attr, val)
+
+
+###
+# The actual Mastodon class
+###
 class Mastodon:
     """
     Thorough and easy to use Mastodon
@@ -3040,8 +3057,8 @@ class Mastodon:
         Returns the next page or None if no further data is available.
         """
         if isinstance(previous_page, list) and len(previous_page) != 0:
-            if hasattr(previous_page[-1], '_pagination_next'):
-                params = copy.deepcopy(previous_page[-1]._pagination_next)
+            if hasattr(previous_page, '_pagination_next'):
+                params = copy.deepcopy(previous_page._pagination_next)
             else:
                 return None
         else:
@@ -3064,8 +3081,8 @@ class Mastodon:
         Returns the previous page or None if no further data is available.
         """
         if isinstance(next_page, list) and len(next_page) != 0:
-            if hasattr(next_page[0], '_pagination_prev'):
-                params = copy.deepcopy(next_page[0]._pagination_prev)
+            if hasattr(next_page, '_pagination_prev'):
+                params = copy.deepcopy(next_page._pagination_prev)
             else:
                 return None
         else:
@@ -3443,6 +3460,7 @@ class Mastodon:
             if isinstance(response, list) and \
                     'Link' in response_object.headers and \
                     response_object.headers['Link'] != "":
+                response = AttribAccessList(response)
                 tmp_urls = requests.utils.parse_header_links(
                     response_object.headers['Link'].rstrip('>').replace('>,<', ',<'))
                 for url in tmp_urls:
@@ -3467,7 +3485,7 @@ class Mastodon:
                                 del next_params['since_id']
                             if "min_id" in next_params:
                                 del next_params['min_id']
-                            response[-1]._pagination_next = next_params
+                            response._pagination_next = next_params
 
                     if url['rel'] == 'prev':
                         # Be paranoid and extract since_id or min_id specifically
@@ -3486,7 +3504,7 @@ class Mastodon:
                                 prev_params['since_id'] = since_id
                             if "max_id" in prev_params:
                                 del prev_params['max_id']
-                            response[0]._pagination_prev = prev_params
+                            response._pagination_prev = prev_params
                             
                         # New and fantastico (post-2.6.0): min_id pagination
                         matchgroups = re.search(r"[?&]min_id=([^&]+)", prev_url)
@@ -3501,7 +3519,7 @@ class Mastodon:
                                 prev_params['min_id'] = min_id
                             if "max_id" in prev_params:
                                 del prev_params['max_id']
-                            response[0]._pagination_prev = prev_params
+                            response._pagination_prev = prev_params
 
         return response