projects.mako.cc - wikipedia-api-cdsw/blob - mwclient/listing.py

   1 import client
   2 import page
   3
   4
   5 class List(object):
   6
   7     def __init__(self, site, list_name, prefix, limit=None, return_values=None, max_items=None, *args, **kwargs):
   8         # NOTE: Fix limit
   9         self.site = site
  10         self.list_name = list_name
  11         self.generator = 'list'
  12         self.prefix = prefix
  13
  14         kwargs.update(args)
  15         self.args = kwargs
  16
  17         if limit is None:
  18             limit = site.api_limit
  19         self.args[self.prefix + 'limit'] = str(limit)
  20         if 'continue' not in self.args:
  21             self.args['continue'] = ''
  22
  23         self.count = 0
  24         self.max_items = max_items
  25
  26         self._iter = iter(xrange(0))
  27
  28         self.last = False
  29         self.result_member = list_name
  30         self.return_values = return_values
  31
  32     def __iter__(self):
  33         return self
  34
  35     def next(self, full=False):
  36         if self.max_items is not None:
  37             if self.count >= self.max_items:
  38                 raise StopIteration
  39         try:
  40             item = self._iter.next()
  41             self.count += 1
  42             if 'timestamp' in item:
  43                 item['timestamp'] = client.parse_timestamp(item['timestamp'])
  44             if full:
  45                 return item
  46
  47             if type(self.return_values) is tuple:
  48                 return tuple((item[i] for i in self.return_values))
  49             elif self.return_values is None:
  50                 return item
  51             else:
  52                 return item[self.return_values]
  53
  54         except StopIteration:
  55             if self.last:
  56                 raise StopIteration
  57             self.load_chunk()
  58             return List.next(self, full=full)
  59
  60     def load_chunk(self):
  61         data = self.site.api('query', (self.generator, self.list_name), *[(str(k), v) for k, v in self.args.iteritems()])
  62         if not data:
  63             # Non existent page
  64             raise StopIteration
  65         self.set_iter(data)
  66
  67         if data.get('continue'):
  68             # New style continuation, added in MediaWiki 1.21
  69             self.args.update(data['continue'])
  70
  71         elif self.list_name in data.get('query-continue', ()):
  72             # Old style continuation
  73             self.args.update(data['query-continue'][self.list_name])
  74
  75         else:
  76             self.last = True
  77
  78     def set_iter(self, data):
  79         if self.result_member not in data['query']:
  80             self._iter = iter(xrange(0))
  81         elif type(data['query'][self.result_member]) is list:
  82             self._iter = iter(data['query'][self.result_member])
  83         else:
  84             self._iter = data['query'][self.result_member].itervalues()
  85
  86     def __repr__(self):
  87         return "<List object '%s' for %s>" % (self.list_name, self.site)
  88
  89     @staticmethod
  90     def generate_kwargs(_prefix, *args, **kwargs):
  91         kwargs.update(args)
  92         for key, value in kwargs.iteritems():
  93             if value is not None and value is not False:
  94                 yield _prefix + key, value
  95
  96     @staticmethod
  97     def get_prefix(prefix, generator=False):
  98         if generator:
  99             return 'g' + prefix
 100         else:
 101             return prefix
 102
 103     @staticmethod
 104     def get_list(generator=False):
 105         if generator:
 106             return GeneratorList
 107         else:
 108             return List
 109
 110
 111 class GeneratorList(List):
 112
 113     def __init__(self, site, list_name, prefix, *args, **kwargs):
 114         List.__init__(self, site, list_name, prefix, *args, **kwargs)
 115
 116         self.args['g' + self.prefix + 'limit'] = self.args[self.prefix + 'limit']
 117         del self.args[self.prefix + 'limit']
 118         self.generator = 'generator'
 119
 120         self.args['prop'] = 'info|imageinfo'
 121         self.args['inprop'] = 'protection'
 122
 123         self.result_member = 'pages'
 124
 125         self.page_class = page.Page
 126
 127     def next(self):
 128         info = List.next(self, full=True)
 129         if info['ns'] == 14:
 130             return Category(self.site, u'', info)
 131         if info['ns'] == 6:
 132             return page.Image(self.site, u'', info)
 133         return page.Page(self.site, u'', info)
 134
 135     def load_chunk(self):
 136         # Put this here so that the constructor does not fail
 137         # on uninitialized sites
 138         self.args['iiprop'] = 'timestamp|user|comment|url|size|sha1|metadata|archivename'
 139         return List.load_chunk(self)
 140
 141
 142 class Category(page.Page, GeneratorList):
 143
 144     def __init__(self, site, name, info=None, namespace=None):
 145         page.Page.__init__(self, site, name, info)
 146         kwargs = {}
 147         kwargs['gcmtitle'] = self.name
 148         if namespace:
 149             kwargs['gcmnamespace'] = namespace
 150         GeneratorList.__init__(self, site, 'categorymembers', 'cm', **kwargs)
 151
 152     def __repr__(self):
 153         return "<Category object '%s' for %s>" % (self.name.encode('utf-8'), self.site)
 154
 155     def members(self, prop='ids|title', namespace=None, sort='sortkey',
 156                 dir='asc', start=None, end=None, generator=True):
 157         prefix = self.get_prefix('cm', generator)
 158         kwargs = dict(self.generate_kwargs(prefix, prop=prop, namespace=namespace,
 159                                            sort=sort, dir=dir, start=start, end=end, title=self.name))
 160         return self.get_list(generator)(self.site, 'categorymembers', 'cm', **kwargs)
 161
 162
 163 class PageList(GeneratorList):
 164
 165     def __init__(self, site, prefix=None, start=None, namespace=0, redirects='all'):
 166         self.namespace = namespace
 167
 168         kwargs = {}
 169         if prefix:
 170             kwargs['apprefix'] = prefix
 171         if start:
 172             kwargs['apfrom'] = start
 173
 174         GeneratorList.__init__(self, site, 'allpages', 'ap',
 175                                apnamespace=str(namespace), apfilterredir=redirects, **kwargs)
 176
 177     def __getitem__(self, name):
 178         return self.get(name, None)
 179
 180     def get(self, name, info=()):
 181         if self.namespace == 14:
 182             return Category(self.site, self.site.namespaces[14] + ':' + name, info)
 183         elif self.namespace == 6:
 184             return page.Image(self.site, self.site.namespaces[6] + ':' + name, info)
 185         elif self.namespace != 0:
 186             return page.Page(self.site, self.site.namespaces[self.namespace] + ':' + name, info)
 187         else:
 188             # Guessing page class
 189             if type(name) is not int:
 190                 namespace = self.guess_namespace(name)
 191                 if namespace == 14:
 192                     return Category(self.site, name, info)
 193                 elif namespace == 6:
 194                     return page.Image(self.site, name, info)
 195             return page.Page(self.site, name, info)
 196
 197     def guess_namespace(self, name):
 198         normal_name = page.Page.normalize_title(name)
 199         for ns in self.site.namespaces:
 200             if ns == 0:
 201                 continue
 202             if name.startswith(u'%s:' % self.site.namespaces[ns].replace(' ', '_')):
 203                 return ns
 204             elif ns in self.site.default_namespaces:
 205                 if name.startswith(u'%s:' % self.site.default_namespaces[ns].replace(' ', '_')):
 206                     return ns
 207         return 0
 208
 209
 210 class PageProperty(List):
 211
 212     def __init__(self, page, prop, prefix, *args, **kwargs):
 213         List.__init__(self, page.site, prop, prefix, titles=page.name, *args, **kwargs)
 214         self.page = page
 215         self.generator = 'prop'
 216
 217     def set_iter(self, data):
 218         for page in data['query']['pages'].itervalues():
 219             if page['title'] == self.page.name:
 220                 self._iter = iter(page.get(self.list_name, ()))
 221                 return
 222         raise StopIteration
 223
 224
 225 class PagePropertyGenerator(GeneratorList):
 226
 227     def __init__(self, page, prop, prefix, *args, **kwargs):
 228         GeneratorList.__init__(self, page.site, prop, prefix, titles=page.name, *args, **kwargs)
 229         self.page = page
 230
 231
 232 class RevisionsIterator(PageProperty):
 233
 234     def load_chunk(self):
 235         if 'rvstartid' in self.args and 'rvstart' in self.args:
 236             del self.args['rvstart']
 237         return PageProperty.load_chunk(self)