From a42db083e3c1a30e8d9364ed5fe802f7ff747903 Mon Sep 17 00:00:00 2001 From: Harmon Date: Sun, 25 Aug 2019 04:22:05 -0500 Subject: [PATCH] Add DMCursorIterator --- tweepy/binder.py | 5 ++++- tweepy/cursor.py | 24 ++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/tweepy/binder.py b/tweepy/binder.py index 7d352a6..ae4ece5 100644 --- a/tweepy/binder.py +++ b/tweepy/binder.py @@ -253,7 +253,10 @@ def bind_api(**config): # Set pagination mode if 'cursor' in APIMethod.allowed_param: - _call.pagination_mode = 'cursor' + if APIMethod.payload_type == 'direct_message': + _call.pagination_mode = 'dm_cursor' + else: + _call.pagination_mode = 'cursor' elif 'max_id' in APIMethod.allowed_param: if 'since_id' in APIMethod.allowed_param: _call.pagination_mode = 'id' diff --git a/tweepy/cursor.py b/tweepy/cursor.py index 1803c37..c928208 100644 --- a/tweepy/cursor.py +++ b/tweepy/cursor.py @@ -13,6 +13,8 @@ class Cursor(object): if hasattr(method, 'pagination_mode'): if method.pagination_mode == 'cursor': self.iterator = CursorIterator(method, *args, **kwargs) + elif method.pagination_mode == 'dm_cursor': + self.iterator = DMCursorIterator(method, *args, **kwargs) elif method.pagination_mode == 'id': self.iterator = IdIterator(method, *args, **kwargs) elif method.pagination_mode == 'page': @@ -87,6 +89,28 @@ class CursorIterator(BaseIterator): return data +class DMCursorIterator(BaseIterator): + + def __init__(self, method, *args, **kwargs): + BaseIterator.__init__(self, method, *args, **kwargs) + self.next_cursor = self.kwargs.pop('cursor', None) + self.page_count = 0 + + def next(self): + if self.next_cursor == -1 or (self.limit and self.page_count == self.limit): + raise StopIteration + data = self.method(cursor=self.next_cursor, *self.args, **self.kwargs) + self.page_count += 1 + if isinstance(data, tuple): + data, self.next_cursor = data + else: + self.next_cursor = -1 + return data + + def prev(self): + raise TweepError('This method does not allow backwards pagination') + + class IdIterator(BaseIterator): def __init__(self, method, *args, **kwargs): -- 2.25.1