Add pagination for Client

diff --git a/tweepy/__init__.py b/tweepy/__init__.py
index 123cbec3cf8de5c5eeb5f3daf960b0d2a1b2cfd6..e1a092ced2981e1e83905017628a3db9b45cd077 100644 (file)
--- a/tweepy/__init__.py
+++ b/tweepy/__init__.py
@@ -17,6 +17,7 @@ from tweepy.cursor import Cursor
 from tweepy.error import RateLimitError, TweepError
 from tweepy.media import Media
 from tweepy.models import DirectMessage, Friendship, ModelFactory, SavedSearch, SearchResults, Status, User
+from tweepy.pagination import Paginator
 from tweepy.place import Place
 from tweepy.poll import Poll
 from tweepy.streaming import Stream

diff --git a/tweepy/pagination.py b/tweepy/pagination.py
new file mode 100644 (file)
index 0000000..68d36fb
--- /dev/null
+++ b/tweepy/pagination.py
@@ -0,0 +1,80 @@
+# Tweepy
+# Copyright 2009-2021 Joshua Roesslein
+# See LICENSE for details.
+
+from math import inf
+
+
+class Paginator:
+
+    def __init__(self, method, *args, **kwargs):
+        self.method = method
+        self.args = args
+        self.kwargs = kwargs
+
+    def __iter__(self):
+        return PaginationIterator(self.method, *self.args, **self.kwargs)
+
+    def __reversed__(self):
+        return PaginationIterator(self.method, *self.args, reverse=True,
+                                  **self.kwargs)
+
+    def flatten(self, limit=inf):
+        if limit <= 0:
+            return
+
+        count = 0
+        for response in PaginationIterator(self.method, *self.args,
+                                           **self.kwargs):
+            for data in response.data:
+                yield data
+                count += 1
+                if count == limit:
+                    return
+
+
+class PaginationIterator:
+
+    def __init__(self, method, *args, limit=inf, pagination_token=None,
+                 reverse=False, **kwargs):
+        self.method = method
+        self.args = args
+        self.limit = limit
+        self.kwargs = kwargs
+        self.reverse = reverse
+
+        if reverse:
+            self.previous_token = pagination_token
+            self.next_token = None
+        else:
+            self.previous_token = None
+            self.next_token = pagination_token
+
+        self.count = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.reverse:
+            pagination_token = self.previous_token
+        else:
+            pagination_token = self.next_token
+
+        if self.count >= self.limit or self.count and pagination_token is None:
+            raise StopIteration
+
+        # https://twittercommunity.com/t/why-does-timeline-use-pagination-token-while-search-uses-next-token/150963
+        if self.method.__name__ in ("search_all_tweets",
+                                    "search_recent_tweets"):
+            self.kwargs["next_token"] = pagination_token
+        else:
+            self.kwargs["pagination_token"] = pagination_token
+
+        response = self.method(*self.args, **self.kwargs)
+
+        self.previous_token = response.meta.get("previous_token")
+        self.next_token = response.meta.get("next_token")
+        self.count += 1
+
+        return response