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

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