Handle different method return types in flatten methods for pagination

tweepy · Oct 27, 2022 · 381bf91 · 381bf91
1 parent e31be15
commit 381bf91
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 14 deletions.
diff --git a/tweepy/asynchronous/pagination.py b/tweepy/asynchronous/pagination.py
@@ -60,12 +60,21 @@ async def flatten(self, limit=inf):
         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
@@ -56,14 +56,23 @@ def flatten(self, limit=inf):
             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: