projects.mako.cc - twitter-api-cdsw-solutions/blob - tweepy/cursor.py

   1 # Tweepy
   2 # Copyright 2009-2010 Joshua Roesslein
   3 # See LICENSE for details.
   4
   5 from tweepy.error import TweepError
   6
   7 class Cursor(object):
   8     """Pagination helper class"""
   9
  10     def __init__(self, method, *args, **kargs):
  11         if hasattr(method, 'pagination_mode'):
  12             if method.pagination_mode == 'cursor':
  13                 self.iterator = CursorIterator(method, args, kargs)
  14             elif method.pagination_mode == 'id':
  15                 self.iterator = IdIterator(method, args, kargs)
  16             elif method.pagination_mode == 'page':
  17                 self.iterator = PageIterator(method, args, kargs)
  18             else:
  19                 raise TweepError('Invalid pagination mode.')
  20         else:
  21             raise TweepError('This method does not perform pagination')
  22
  23     def pages(self, limit=0):
  24         """Return iterator for pages"""
  25         if limit > 0:
  26             self.iterator.limit = limit
  27         return self.iterator
  28
  29     def items(self, limit=0):
  30         """Return iterator for items in each page"""
  31         i = ItemIterator(self.iterator)
  32         i.limit = limit
  33         return i
  34
  35 class BaseIterator(object):
  36
  37     def __init__(self, method, args, kargs):
  38         self.method = method
  39         self.args = args
  40         self.kargs = kargs
  41         self.limit = 0
  42
  43     def next(self):
  44         raise NotImplementedError
  45
  46     def prev(self):
  47         raise NotImplementedError
  48
  49     def __iter__(self):
  50         return self
  51
  52 class CursorIterator(BaseIterator):
  53
  54     def __init__(self, method, args, kargs):
  55         BaseIterator.__init__(self, method, args, kargs)
  56         start_cursor = kargs.pop('cursor', None)
  57         self.next_cursor = start_cursor or -1
  58         self.prev_cursor = start_cursor or 0
  59         self.count = 0
  60
  61     def next(self):
  62         if self.next_cursor == 0 or (self.limit and self.count == self.limit):
  63             raise StopIteration
  64         data, cursors = self.method(
  65                 cursor=self.next_cursor, *self.args, **self.kargs
  66         )
  67         self.prev_cursor, self.next_cursor = cursors
  68         if len(data) == 0:
  69             raise StopIteration
  70         self.count += 1
  71         return data
  72
  73     def prev(self):
  74         if self.prev_cursor == 0:
  75             raise TweepError('Can not page back more, at first page')
  76         data, self.next_cursor, self.prev_cursor = self.method(
  77                 cursor=self.prev_cursor, *self.args, **self.kargs
  78         )
  79         self.count -= 1
  80         return data
  81
  82 class IdIterator(BaseIterator):
  83
  84     def __init__(self, method, args, kargs):
  85         BaseIterator.__init__(self, method, args, kargs)
  86         self.max_id = kargs.get('max_id')
  87         self.since_id = kargs.get('since_id')
  88         self.count = 0
  89
  90     def next(self):
  91         """Fetch a set of items with IDs less than current set."""
  92         if self.limit and self.limit == self.count:
  93             raise StopIteration
  94
  95         # max_id is inclusive so decrement by one
  96         # to avoid requesting duplicate items.
  97         max_id = self.since_id - 1 if self.max_id else None
  98         data = self.method(max_id = max_id, *self.args, **self.kargs)
  99         if len(data) == 0:
 100             raise StopIteration
 101         self.max_id = data.max_id
 102         self.since_id = data.since_id
 103         self.count += 1
 104         return data
 105
 106     def prev(self):
 107         """Fetch a set of items with IDs greater than current set."""
 108         if self.limit and self.limit == self.count:
 109             raise StopIteration
 110
 111         since_id = self.max_id
 112         data = self.method(since_id = since_id, *self.args, **self.kargs)
 113         if len(data) == 0:
 114             raise StopIteration
 115         self.max_id = data.max_id
 116         self.since_id = data.since_id
 117         self.count += 1
 118         return data
 119
 120 class PageIterator(BaseIterator):
 121
 122     def __init__(self, method, args, kargs):
 123         BaseIterator.__init__(self, method, args, kargs)
 124         self.current_page = 0
 125
 126     def next(self):
 127         self.current_page += 1
 128         items = self.method(page=self.current_page, *self.args, **self.kargs)
 129         if len(items) == 0 or (self.limit > 0 and self.current_page > self.limit):
 130             raise StopIteration
 131         return items
 132
 133     def prev(self):
 134         if (self.current_page == 1):
 135             raise TweepError('Can not page back more, at first page')
 136         self.current_page -= 1
 137         return self.method(page=self.current_page, *self.args, **self.kargs)
 138
 139 class ItemIterator(BaseIterator):
 140
 141     def __init__(self, page_iterator):
 142         self.page_iterator = page_iterator
 143         self.limit = 0
 144         self.current_page = None
 145         self.page_index = -1
 146         self.count = 0
 147
 148     def next(self):
 149         if self.limit > 0 and self.count == self.limit:
 150             raise StopIteration
 151         if self.current_page is None or self.page_index == len(self.current_page) - 1:
 152             # Reached end of current page, get the next page...
 153             self.current_page = self.page_iterator.next()
 154             self.page_index = -1
 155         self.page_index += 1
 156         self.count += 1
 157         return self.current_page[self.page_index]
 158
 159     def prev(self):
 160         if self.current_page is None:
 161             raise TweepError('Can not go back more, at first page')
 162         if self.page_index == 0:
 163             # At the beginning of the current page, move to next...
 164             self.current_page = self.page_iterator.prev()
 165             self.page_index = len(self.current_page)
 166             if self.page_index == 0:
 167                 raise TweepError('No more items')
 168         self.page_index -= 1
 169         self.count -= 1
 170         return self.current_page[self.page_index]
 171