From: Harmon <Harmon758@gmail.com>
Date: Thu, 27 Oct 2022 17:33:49 +0000 (-0500)
Subject: Handle different method return types in flatten methods for pagination
X-Git-Url: https://vcs.fsf.org/?a=commitdiff_plain;h=381bf9182700fa0c51df00b5b0f3e19c3d04dac0;p=tweepy.git

Handle different method return types in flatten methods for pagination
---

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: