From 381bf9182700fa0c51df00b5b0f3e19c3d04dac0 Mon Sep 17 00:00:00 2001 From: Harmon Date: Thu, 27 Oct 2022 12:33:49 -0500 Subject: [PATCH] Handle different method return types in flatten methods for pagination --- tweepy/asynchronous/pagination.py | 21 +++++++++++++++------ tweepy/pagination.py | 25 +++++++++++++++++-------- 2 files changed, 32 insertions(+), 14 deletions(-) diff --git a/tweepy/asynchronous/pagination.py b/tweepy/asynchronous/pagination.py index ccbd34f..6743e0b 100644 --- a/tweepy/asynchronous/pagination.py +++ b/tweepy/asynchronous/pagination.py @@ -60,12 +60,21 @@ class AsyncPaginator: 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 + if isinstance(response, Response): + response_data = response.data or [] + elif isinstance(response, dict): + response_data = response.get("data", []) + else: + raise RuntimeError( + "AsyncPaginator.flatten does not support the " + f"{type(response)} return type for " + f"{self.method.__qualname__}" + ) + for data in response_data: + yield data + count += 1 + if count == limit: + return class AsyncPaginationIterator: diff --git a/tweepy/pagination.py b/tweepy/pagination.py index 91abc8d..cdf429d 100644 --- a/tweepy/pagination.py +++ b/tweepy/pagination.py @@ -56,14 +56,23 @@ class Paginator: return count = 0 - for response in PaginationIterator(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 + for response in PaginationIterator( + self.method, *self.args, **self.kwargs + ): + if isinstance(response, Response): + response_data = response.data or [] + elif isinstance(response, dict): + response_data = response.get("data", []) + else: + raise RuntimeError( + f"Paginator.flatten does not support the {type(response)} " + f"return type for {self.method.__qualname__}" + ) + for data in response_data: + yield data + count += 1 + if count == limit: + return class PaginationIterator: -- 2.25.1