Handle different method return types in flatten methods for pagination
authorHarmon <Harmon758@gmail.com>
Thu, 27 Oct 2022 17:33:49 +0000 (12:33 -0500)
committerHarmon <Harmon758@gmail.com>
Thu, 27 Oct 2022 17:50:29 +0000 (12:50 -0500)
tweepy/asynchronous/pagination.py
tweepy/pagination.py

index ccbd34f3e90bcfddccd50a7eff7205aaab9cf9df..6743e0b61d48bffe9a913096df4e783a2a1c0632 100644 (file)
@@ -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:
index 91abc8d6774a0fc53753145ca0ae8abfc7dec52c..cdf429dae37c84bf0696504e71536e9c8e3f7bd7 100644 (file)
@@ -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: