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 __future__ import print_function
   6
   7 from tweepy.error import TweepError
   8 from tweepy.parsers import ModelParser, RawParser
   9
  10
  11 class Cursor(object):
  12     """Pagination helper class"""
  13
  14     def __init__(self, method, *args, **kargs):
  15         if hasattr(method, 'pagination_mode'):
  16             if method.pagination_mode == 'cursor':
  17                 self.iterator = CursorIterator(method, args, kargs)
  18             elif method.pagination_mode == 'id':
  19                 self.iterator = IdIterator(method, args, kargs)
  20             elif method.pagination_mode == 'page':
  21                 self.iterator = PageIterator(method, args, kargs)
  22             else:
  23                 raise TweepError('Invalid pagination mode.')
  24         else:
  25             raise TweepError('This method does not perform pagination')
  26
  27     def pages(self, limit=0):
  28         """Return iterator for pages"""
  29         if limit > 0:
  30             self.iterator.limit = limit
  31         return self.iterator
  32
  33     def items(self, limit=0):
  34         """Return iterator for items in each page"""
  35         i = ItemIterator(self.iterator)
  36         i.limit = limit
  37         return i
  38
  39
  40 class BaseIterator(object):
  41
  42     def __init__(self, method, args, kargs):
  43         self.method = method
  44         self.args = args
  45         self.kargs = kargs
  46         self.limit = 0
  47
  48     def __next__(self):
  49         return self.next()
  50
  51     def next(self):
  52         raise NotImplementedError
  53
  54     def prev(self):
  55         raise NotImplementedError
  56
  57     def __iter__(self):
  58         return self
  59
  60
  61 class CursorIterator(BaseIterator):
  62
  63     def __init__(self, method, args, kargs):
  64         BaseIterator.__init__(self, method, args, kargs)
  65         start_cursor = kargs.pop('cursor', None)
  66         self.next_cursor = start_cursor or -1
  67         self.prev_cursor = start_cursor or 0
  68         self.num_tweets = 0
  69
  70     def next(self):
  71         if self.next_cursor == 0 or (self.limit and self.num_tweets == self.limit):
  72             raise StopIteration
  73         data, cursors = self.method(cursor=self.next_cursor,
  74                                     *self.args,
  75                                     **self.kargs)
  76         self.prev_cursor, self.next_cursor = cursors
  77         if len(data) == 0:
  78             raise StopIteration
  79         self.num_tweets += 1
  80         return data
  81
  82     def prev(self):
  83         if self.prev_cursor == 0:
  84             raise TweepError('Can not page back more, at first page')
  85         data, self.next_cursor, self.prev_cursor = self.method(cursor=self.prev_cursor,
  86                                                                *self.args,
  87                                                                **self.kargs)
  88         self.num_tweets -= 1
  89         return data
  90
  91
  92 class IdIterator(BaseIterator):
  93
  94     def __init__(self, method, args, kargs):
  95         BaseIterator.__init__(self, method, args, kargs)
  96         self.max_id = kargs.pop('max_id', None)
  97         self.num_tweets = 0
  98         self.results = []
  99         self.model_results = []
 100         self.index = 0
 101
 102     def next(self):
 103         """Fetch a set of items with IDs less than current set."""
 104         if self.limit and self.limit == self.num_tweets:
 105             raise StopIteration
 106
 107         if self.index >= len(self.results) - 1:
 108             data = self.method(max_id=self.max_id, parser=RawParser(), *self.args, **self.kargs)
 109
 110             if hasattr(self.method, '__self__'):
 111                 old_parser = self.method.__self__.parser
 112                 # Hack for models which expect ModelParser to be set
 113                 self.method.__self__.parser = ModelParser()
 114
 115             # This is a special invocation that returns the underlying
 116             # APIMethod class
 117             model = ModelParser().parse(self.method(create=True), data)
 118             if hasattr(self.method, '__self__'):
 119                 self.method.__self__.parser = old_parser
 120                 result = self.method.__self__.parser.parse(self.method(create=True), data)
 121             else:
 122                 result = model
 123
 124             if len(self.results) != 0:
 125                 self.index += 1
 126             self.results.append(result)
 127             self.model_results.append(model)
 128         else:
 129             self.index += 1
 130             result = self.results[self.index]
 131             model = self.model_results[self.index]
 132
 133         if len(result) == 0:
 134             raise StopIteration
 135         # TODO: Make this not dependant on the parser making max_id and
 136         # since_id available
 137         self.max_id = model.max_id
 138         self.num_tweets += 1
 139         return result
 140
 141     def prev(self):
 142         """Fetch a set of items with IDs greater than current set."""
 143         if self.limit and self.limit == self.num_tweets:
 144             raise StopIteration
 145
 146         self.index -= 1
 147         if self.index < 0:
 148             # There's no way to fetch a set of tweets directly 'above' the
 149             # current set
 150             raise StopIteration
 151
 152         data = self.results[self.index]
 153         self.max_id = self.model_results[self.index].max_id
 154         self.num_tweets += 1
 155         return data
 156
 157
 158 class PageIterator(BaseIterator):
 159
 160     def __init__(self, method, args, kargs):
 161         BaseIterator.__init__(self, method, args, kargs)
 162         self.current_page = 0
 163
 164     def next(self):
 165         if self.limit > 0:
 166             if self.current_page > self.limit:
 167                 raise StopIteration
 168
 169         items = self.method(page=self.current_page, *self.args, **self.kargs)
 170         if len(items) == 0:
 171             raise StopIteration
 172         self.current_page += 1
 173         return items
 174
 175     def prev(self):
 176         if self.current_page == 1:
 177             raise TweepError('Can not page back more, at first page')
 178         self.current_page -= 1
 179         return self.method(page=self.current_page, *self.args, **self.kargs)
 180
 181
 182 class ItemIterator(BaseIterator):
 183
 184     def __init__(self, page_iterator):
 185         self.page_iterator = page_iterator
 186         self.limit = 0
 187         self.current_page = None
 188         self.page_index = -1
 189         self.num_tweets = 0
 190
 191     def next(self):
 192         if self.limit > 0:
 193             if self.num_tweets == self.limit:
 194                 raise StopIteration
 195         if self.current_page is None or self.page_index == len(self.current_page) - 1:
 196             # Reached end of current page, get the next page...
 197             self.current_page = self.page_iterator.next()
 198             self.page_index = -1
 199         self.page_index += 1
 200         self.num_tweets += 1
 201         return self.current_page[self.page_index]
 202
 203     def prev(self):
 204         if self.current_page is None:
 205             raise TweepError('Can not go back more, at first page')
 206         if self.page_index == 0:
 207             # At the beginning of the current page, move to next...
 208             self.current_page = self.page_iterator.prev()
 209             self.page_index = len(self.current_page)
 210             if self.page_index == 0:
 211                 raise TweepError('No more items')
 212         self.page_index -= 1
 213         self.num_tweets -= 1
 214         return self.current_page[self.page_index]