From 33e444a9d13d53ea024ddb3c9da30158a39ea4f6 Mon Sep 17 00:00:00 2001 From: Harmon Date: Mon, 10 Oct 2022 00:25:14 -0500 Subject: [PATCH] Add pagination for AsyncClient --- docs/v2_pagination.rst | 3 + tweepy/asynchronous/__init__.py | 1 + tweepy/asynchronous/pagination.py | 108 ++++++++++++++++++++++++++++++ 3 files changed, 112 insertions(+) create mode 100644 tweepy/asynchronous/pagination.py diff --git a/docs/v2_pagination.rst b/docs/v2_pagination.rst index 8b3db43..f8d5813 100644 --- a/docs/v2_pagination.rst +++ b/docs/v2_pagination.rst @@ -24,3 +24,6 @@ Pagination for tweet in tweepy.Paginator(client.search_recent_tweets, "Tweepy", max_results=100).flatten(limit=250): print(tweet.id) + +.. autoclass:: tweepy.asynchronous.AsyncPaginator + :members: diff --git a/tweepy/asynchronous/__init__.py b/tweepy/asynchronous/__init__.py index d5dc1a8..3bd4dc2 100644 --- a/tweepy/asynchronous/__init__.py +++ b/tweepy/asynchronous/__init__.py @@ -21,3 +21,4 @@ except ModuleNotFoundError: from tweepy.asynchronous.streaming import AsyncStream, AsyncStreamingClient from tweepy.asynchronous.client import AsyncClient +from tweepy.asynchronous.pagination import AsyncPaginator diff --git a/tweepy/asynchronous/pagination.py b/tweepy/asynchronous/pagination.py new file mode 100644 index 0000000..caf7cfa --- /dev/null +++ b/tweepy/asynchronous/pagination.py @@ -0,0 +1,108 @@ +# Tweepy +# Copyright 2009-2022 Joshua Roesslein +# See LICENSE for details. + +from math import inf + + +class AsyncPaginator: + """:class:`AsyncPaginator` can be used to paginate for any + :class:`AsyncClient` methods that support pagination + + .. versionadded:: 4.11 + + Parameters + ---------- + method + :class:`AsyncClient` method to paginate for + args + Positional arguments to pass to ``method`` + kwargs + Keyword arguments to pass to ``method`` + """ + + def __init__(self, method, *args, **kwargs): + self.method = method + self.args = args + self.kwargs = kwargs + + def __aiter__(self): + return AsyncPaginationIterator(self.method, *self.args, **self.kwargs) + + def __reversed__(self): + return AsyncPaginationIterator( + self.method, *self.args, reverse=True, **self.kwargs + ) + + async def flatten(self, limit=inf): + """Flatten paginated data + + Parameters + ---------- + limit + Maximum number of results to yield + """ + if limit <= 0: + return + + count = 0 + async for response in AsyncPaginationIterator( + self.method, *self.args, **self.kwargs + ): + if response.data is not None: + for data in response.data: + yield data + count += 1 + if count == limit: + return + + +class AsyncPaginationIterator: + + 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 __aiter__(self): + return self + + async def __anext__(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 StopAsyncIteration + + # 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", + "get_all_tweets_count" + ): + self.kwargs["next_token"] = pagination_token + else: + self.kwargs["pagination_token"] = pagination_token + + response = await 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 -- 2.25.1