c628662432f371383cf64cc9069715fcc6357d98
[wikipedia-api-cdsw] / mwclient / client.py
1 __ver__ = '0.6.6'
2
3 import urllib
4 import urlparse
5 import time
6 import random
7 import sys
8 import weakref
9 import socket
10 import base64
11
12 try:
13     import json
14 except ImportError:
15     import simplejson as json
16 import http
17 import upload
18
19 import errors
20 import listing
21 import page
22 import compatibility
23
24 try:
25     import gzip
26 except ImportError:
27     gzip = None
28 try:
29     from cStringIO import StringIO
30 except ImportError:
31     from StringIO import StringIO
32
33
34 def parse_timestamp(t):
35     if t == '0000-00-00T00:00:00Z':
36         return (0, 0, 0, 0, 0, 0, 0, 0)
37     return time.strptime(t, '%Y-%m-%dT%H:%M:%SZ')
38
39
40 class WaitToken(object):
41
42     def __init__(self):
43         self.id = '%x' % random.randint(0, sys.maxint)
44
45     def __hash__(self):
46         return hash(self.id)
47
48
49 class Site(object):
50     api_limit = 500
51
52     def __init__(self, host, path='/w/', ext='.php', pool=None, retry_timeout=30,
53                  max_retries=25, wait_callback=lambda *x: None, clients_useragent=None,
54                  max_lag=3, compress=True, force_login=True, do_init=True, httpauth=None):
55         # Setup member variables
56         self.host = host
57         self.path = path
58         self.ext = ext
59         self.credentials = None
60         self.compress = compress
61         self.httpauth = httpauth
62         self.retry_timeout = retry_timeout
63         self.max_retries = max_retries
64         self.wait_callback = wait_callback
65         self.max_lag = str(max_lag)
66         self.force_login = force_login
67
68         # The token string => token object mapping
69         self.wait_tokens = weakref.WeakKeyDictionary()
70
71         # Site properties
72         self.blocked = False    # Whether current user is blocked
73         self.hasmsg = False  # Whether current user has new messages
74         self.groups = []    # Groups current user belongs to
75         self.rights = []    # Rights current user has
76         self.tokens = {}    # Edit tokens of the current user
77         self.version = None
78
79         self.namespaces = self.default_namespaces
80         self.writeapi = False
81
82         # Setup connection
83         if pool is None:
84             self.connection = http.HTTPPool(clients_useragent)
85         else:
86             self.connection = pool
87
88         # Page generators
89         self.pages = listing.PageList(self)
90         self.categories = listing.PageList(self, namespace=14)
91         self.images = listing.PageList(self, namespace=6)
92
93         # Compat page generators
94         self.Pages = self.pages
95         self.Categories = self.categories
96         self.Images = self.images
97
98         # Initialization status
99         self.initialized = False
100
101         if do_init:
102             try:
103                 self.site_init()
104             except errors.APIError, e:
105                 # Private wiki, do init after login
106                 if e[0] not in (u'unknown_action', u'readapidenied'):
107                     raise
108
109     def site_init(self):
110         meta = self.api('query', meta='siteinfo|userinfo',
111                         siprop='general|namespaces', uiprop='groups|rights')
112
113         # Extract site info
114         self.site = meta['query']['general']
115         self.namespaces = dict(((i['id'], i.get('*', '')) for i in meta['query']['namespaces'].itervalues()))
116         self.writeapi = 'writeapi' in self.site
117
118         # Determine version
119         if self.site['generator'].startswith('MediaWiki '):
120             version = self.site['generator'][10:].split('.')
121
122             def split_num(s):
123                 i = 0
124                 while i < len(s):
125                     if s[i] < '0' or s[i] > '9':
126                         break
127                     i += 1
128                 if s[i:]:
129                     return (int(s[:i]), s[i:], )
130                 else:
131                     return (int(s[:i]), )
132             self.version = sum((split_num(s) for s in version), ())
133
134             if len(self.version) < 2:
135                 raise errors.MediaWikiVersionError('Unknown MediaWiki %s' % '.'.join(version))
136         else:
137             raise errors.MediaWikiVersionError('Unknown generator %s' % self.site['generator'])
138         # Require 1.11 until some compatibility issues are fixed
139         self.require(1, 11)
140
141         # User info
142         userinfo = compatibility.userinfo(meta, self.require(1, 12, raise_error=False))
143         self.username = userinfo['name']
144         self.groups = userinfo.get('groups', [])
145         self.rights = userinfo.get('rights', [])
146         self.initialized = True
147
148     default_namespaces = {0: u'', 1: u'Talk', 2: u'User', 3: u'User talk', 4: u'Project', 5: u'Project talk',
149                           6: u'Image', 7: u'Image talk', 8: u'MediaWiki', 9: u'MediaWiki talk', 10: u'Template', 11: u'Template talk',
150                           12: u'Help', 13: u'Help talk', 14: u'Category', 15: u'Category talk', -1: u'Special', -2: u'Media'}
151
152     def __repr__(self):
153         return "<Site object '%s%s'>" % (self.host, self.path)
154
155     def api(self, action, *args, **kwargs):
156         """ An API call. Handles errors and returns dict object. """
157         kwargs.update(args)
158         if action == 'query':
159             if 'meta' in kwargs:
160                 kwargs['meta'] += '|userinfo'
161             else:
162                 kwargs['meta'] = 'userinfo'
163             if 'uiprop' in kwargs:
164                 kwargs['uiprop'] += '|blockinfo|hasmsg'
165             else:
166                 kwargs['uiprop'] = 'blockinfo|hasmsg'
167
168         token = self.wait_token()
169         while True:
170             info = self.raw_api(action, **kwargs)
171             if not info:
172                 info = {}
173             res = self.handle_api_result(info, token=token)
174             if res:
175                 return info
176
177     def handle_api_result(self, info, kwargs=None, token=None):
178         if token is None:
179             token = self.wait_token()
180
181         try:
182             userinfo = compatibility.userinfo(info, self.require(1, 12, raise_error=None))
183         except KeyError:
184             userinfo = ()
185         if 'blockedby' in userinfo:
186             self.blocked = (userinfo['blockedby'], userinfo.get('blockreason', u''))
187         else:
188             self.blocked = False
189         self.hasmsg = 'message' in userinfo
190         self.logged_in = 'anon' not in userinfo
191         if 'error' in info:
192             if info['error']['code'] in (u'internal_api_error_DBConnectionError', ):
193                 self.wait(token)
194                 return False
195             if '*' in info['error']:
196                 raise errors.APIError(info['error']['code'],
197                                       info['error']['info'], info['error']['*'])
198             raise errors.APIError(info['error']['code'],
199                                   info['error']['info'], kwargs)
200         return True
201
202     @staticmethod
203     def _to_str(data):
204         if type(data) is unicode:
205             return data.encode('utf-8')
206         return str(data)
207
208     @staticmethod
209     def _query_string(*args, **kwargs):
210         kwargs.update(args)
211         qs = urllib.urlencode([(k, Site._to_str(v)) for k, v in kwargs.iteritems()
212                                if k != 'wpEditToken'])
213         if 'wpEditToken' in kwargs:
214             qs += '&wpEditToken=' + urllib.quote(Site._to_str(kwargs['wpEditToken']))
215         return qs
216
217     def raw_call(self, script, data):
218         url = self.path + script + self.ext
219         headers = {}
220         if not issubclass(data.__class__, upload.Upload):
221             headers['Content-Type'] = 'application/x-www-form-urlencoded'
222         if self.compress and gzip:
223             headers['Accept-Encoding'] = 'gzip'
224         if self.httpauth is not None:
225             credentials = base64.encodestring('%s:%s' % self.httpauth).replace('\n', '')
226             headers['Authorization'] = 'Basic %s' % credentials
227         token = self.wait_token((script, data))
228         while True:
229             try:
230                 stream = self.connection.post(self.host,
231                                               url, data=data, headers=headers)
232                 if stream.getheader('Content-Encoding') == 'gzip':
233                     # BAD.
234                     seekable_stream = StringIO(stream.read())
235                     stream = gzip.GzipFile(fileobj=seekable_stream)
236                 return stream
237
238             except errors.HTTPStatusError, e:
239                 if e[0] == 503 and e[1].getheader('X-Database-Lag'):
240                     self.wait(token, int(e[1].getheader('Retry-After')))
241                 elif e[0] < 500 or e[0] > 599:
242                     raise
243                 else:
244                     self.wait(token)
245             except errors.HTTPRedirectError:
246                 raise
247             except errors.HTTPError:
248                 self.wait(token)
249             except ValueError:
250                 self.wait(token)
251
252     def raw_api(self, action, *args, **kwargs):
253         """Sends a call to the API."""
254         kwargs['action'] = action
255         kwargs['format'] = 'json'
256         data = self._query_string(*args, **kwargs)
257         json_data = self.raw_call('api', data).read()
258         try:
259             return json.loads(json_data)
260         except ValueError:
261             if json_data.startswith('MediaWiki API is not enabled for this site.'):
262                 raise errors.APIDisabledError
263             raise
264
265     def raw_index(self, action, *args, **kwargs):
266         """Sends a call to index.php rather than the API."""
267         kwargs['action'] = action
268         kwargs['maxlag'] = self.max_lag
269         data = self._query_string(*args, **kwargs)
270         return self.raw_call('index', data).read().decode('utf-8', 'ignore')
271
272     def wait_token(self, args=None):
273         token = WaitToken()
274         self.wait_tokens[token] = (0, args)
275         return token
276
277     def wait(self, token, min_wait=0):
278         retry, args = self.wait_tokens[token]
279         self.wait_tokens[token] = (retry + 1, args)
280         if retry > self.max_retries and self.max_retries != -1:
281             raise errors.MaximumRetriesExceeded(self, token, args)
282         self.wait_callback(self, token, retry, args)
283
284         timeout = self.retry_timeout * retry
285         if timeout < min_wait:
286             timeout = min_wait
287         time.sleep(timeout)
288         return self.wait_tokens[token]
289
290     def require(self, major, minor, revision=None, raise_error=True):
291         if self.version is None:
292             if raise_error is None:
293                 return
294             raise RuntimeError('Site %s has not yet been initialized' % repr(self))
295
296         if revision is None:
297             if self.version[:2] >= (major, minor):
298                 return True
299             elif raise_error:
300                 raise errors.MediaWikiVersionError('Requires version %s.%s, current version is %s.%s'
301                                                    % ((major, minor) + self.version[:2]))
302             else:
303                 return False
304         else:
305             raise NotImplementedError
306
307     # Actions
308     def email(self, user, text, subject, cc=False):
309         """Sends email to a specified user on the wiki."""
310         # TODO: Use api!
311         postdata = {}
312         postdata['wpSubject'] = subject
313         postdata['wpText'] = text
314         if cc:
315             postdata['wpCCMe'] = '1'
316         postdata['wpEditToken'] = self.tokens['edit']
317         postdata['uselang'] = 'en'
318         postdata['title'] = u'Special:Emailuser/' + user
319
320         data = self.raw_index('submit', **postdata)
321         if 'var wgAction = "success";' not in data:
322             if 'This user has not specified a valid e-mail address' in data:
323                 # Dirty hack
324                 raise errors.NoSpecifiedEmailError, user
325             raise errors.EmailError, data
326
327     def login(self, username=None, password=None, cookies=None, domain=None):
328         """Login to the wiki."""
329         if self.initialized:
330             self.require(1, 10)
331
332         if username and password:
333             self.credentials = (username, password, domain)
334         if cookies:
335             if self.host not in self.conn.cookies:
336                 self.conn.cookies[self.host] = http.CookieJar()
337             self.conn.cookies[self.host].update(cookies)
338
339         if self.credentials:
340             wait_token = self.wait_token()
341             kwargs = {
342                 'lgname': self.credentials[0],
343                 'lgpassword': self.credentials[1]
344             }
345             if self.credentials[2]:
346                 kwargs['lgdomain'] = self.credentials[2]
347             while True:
348                 login = self.api('login', **kwargs)
349                 if login['login']['result'] == 'Success':
350                     break
351                 elif login['login']['result'] == 'NeedToken':
352                     kwargs['lgtoken'] = login['login']['token']
353                 elif login['login']['result'] == 'Throttled':
354                     self.wait(wait_token, login['login'].get('wait', 5))
355                 else:
356                     raise errors.LoginError(self, login['login'])
357
358         if self.initialized:
359             info = self.api('query', meta='userinfo', uiprop='groups|rights')
360             userinfo = compatibility.userinfo(info, self.require(1, 12, raise_error=False))
361             self.username = userinfo['name']
362             self.groups = userinfo.get('groups', [])
363             self.rights = userinfo.get('rights', [])
364             self.tokens = {}
365         else:
366             self.site_init()
367
368     def upload(self, file=None, filename=None, description='', ignore=False, file_size=None,
369                url=None, session_key=None, comment=None):
370         """Upload a file to the wiki."""
371         if self.version[:2] < (1, 16):
372             return compatibility.old_upload(self, file=file, filename=filename,
373                                             description=description, ignore=ignore,
374                                             file_size=file_size)
375
376         image = self.Images[filename]
377         if not image.can('upload'):
378             raise errors.InsufficientPermission(filename)
379
380         predata = {}
381
382         if comment is None:
383             predata['comment'] = description
384         else:
385             predata['comment'] = comment
386             predata['text'] = description
387
388         if ignore:
389             predata['ignorewarnings'] = 'true'
390         predata['token'] = image.get_token('edit')
391         predata['action'] = 'upload'
392         predata['format'] = 'json'
393         predata['filename'] = filename
394         if url:
395             predata['url'] = url
396         if session_key:
397             predata['session_key'] = session_key
398
399         if file is None:
400             postdata = self._query_string(predata)
401         else:
402             if type(file) is str:
403                 file_size = len(file)
404                 file = StringIO(file)
405             if file_size is None:
406                 file.seek(0, 2)
407                 file_size = file.tell()
408                 file.seek(0, 0)
409
410             postdata = upload.UploadFile('file', filename, file_size, file, predata)
411
412         wait_token = self.wait_token()
413         while True:
414             try:
415                 data = self.raw_call('api', postdata).read()
416                 info = json.loads(data)
417                 if not info:
418                     info = {}
419                 if self.handle_api_result(info, kwargs=predata):
420                     return info.get('upload', {})
421             except errors.HTTPStatusError, e:
422                 if e[0] == 503 and e[1].getheader('X-Database-Lag'):
423                     self.wait(wait_token, int(e[1].getheader('Retry-After')))
424                 elif e[0] < 500 or e[0] > 599:
425                     raise
426                 else:
427                     self.wait(wait_token)
428             except errors.HTTPError:
429                 self.wait(wait_token)
430             file.seek(0, 0)
431
432     def parse(self, text=None, title=None, page=None):
433         kwargs = {}
434         if text is not None:
435             kwargs['text'] = text
436         if title is not None:
437             kwargs['title'] = title
438         if page is not None:
439             kwargs['page'] = page
440         result = self.api('parse', **kwargs)
441         return result['parse']
442
443     # def block: requires 1.12
444     # def unblock: requires 1.12
445     # def patrol: requires 1.14
446     # def import: requires 1.15
447
448     # Lists
449     def allpages(self, start=None, prefix=None, namespace='0', filterredir='all',
450                  minsize=None, maxsize=None, prtype=None, prlevel=None,
451                  limit=None, dir='ascending', filterlanglinks='all', generator=True):
452         """Retrieve all pages on the wiki as a generator."""
453         self.require(1, 9)
454
455         pfx = listing.List.get_prefix('ap', generator)
456         kwargs = dict(listing.List.generate_kwargs(pfx, ('from', start), prefix=prefix,
457                                                    minsize=minsize, maxsize=maxsize, prtype=prtype, prlevel=prlevel,
458                                                    namespace=namespace, filterredir=filterredir, dir=dir,
459                                                    filterlanglinks=filterlanglinks))
460         return listing.List.get_list(generator)(self, 'allpages', 'ap', limit=limit, return_values='title', **kwargs)
461     # def allimages(self): requires 1.12
462     # TODO!
463
464     def alllinks(self, start=None, prefix=None, unique=False, prop='title',
465                  namespace='0', limit=None, generator=True):
466         """Retrieve a list of all links on the wiki as a generator."""
467         self.require(1, 11)
468
469         pfx = listing.List.get_prefix('al', generator)
470         kwargs = dict(listing.List.generate_kwargs(pfx, ('from', start), prefix=prefix,
471                                                    prop=prop, namespace=namespace))
472         if unique:
473             kwargs[pfx + 'unique'] = '1'
474         return listing.List.get_list(generator)(self, 'alllinks', 'al', limit=limit, return_values='title', **kwargs)
475
476     def allcategories(self, start=None, prefix=None, dir='ascending', limit=None, generator=True):
477         """Retrieve all categories on the wiki as a generator."""
478         self.require(1, 12)
479
480         pfx = listing.List.get_prefix('ac', generator)
481         kwargs = dict(listing.List.generate_kwargs(pfx, ('from', start), prefix=prefix, dir=dir))
482         return listing.List.get_list(generator)(self, 'allcategories', 'ac', limit=limit, **kwargs)
483
484     def allusers(self, start=None, prefix=None, group=None, prop=None, limit=None):
485         """Retrieve all users on the wiki as a generator."""
486         self.require(1, 11)
487
488         kwargs = dict(listing.List.generate_kwargs('au', ('from', start), prefix=prefix,
489                                                    group=group, prop=prop))
490         return listing.List(self, 'allusers', 'au', limit=limit, **kwargs)
491
492     def blocks(self, start=None, end=None, dir='older', ids=None, users=None, limit=None,
493                prop='id|user|by|timestamp|expiry|reason|flags'):
494         """Retrieve blocks as a generator.
495
496         Each block is a dictionary containing:
497         - user: the username or IP address of the user
498         - id: the ID of the block
499         - timestamp: when the block was added
500         - expiry: when the block runs out (infinity for indefinite blocks)
501         - reason: the reason they are blocked
502         - allowusertalk: key is present (empty string) if the user is allowed to edit their user talk page
503         - by: the administrator who blocked the user
504         - nocreate: key is present (empty string) if the user's ability to create accounts has been disabled.
505
506         """
507
508         self.require(1, 12)
509         # TODO: Fix. Fix what?
510         kwargs = dict(listing.List.generate_kwargs('bk', start=start, end=end, dir=dir,
511                                                    users=users, prop=prop))
512         return listing.List(self, 'blocks', 'bk', limit=limit, **kwargs)
513
514     def deletedrevisions(self, start=None, end=None, dir='older', namespace=None,
515                          limit=None, prop='user|comment'):
516         # TODO: Fix
517         self.require(1, 12)
518
519         kwargs = dict(listing.List.generate_kwargs('dr', start=start, end=end, dir=dir,
520                                                    namespace=namespace, prop=prop))
521         return listing.List(self, 'deletedrevs', 'dr', limit=limit, **kwargs)
522
523     def exturlusage(self, query, prop=None, protocol='http', namespace=None, limit=None):
524         """Retrieves list of pages that link to a particular domain or URL as a generator.
525
526         This API call mirrors the Special:LinkSearch function on-wiki.
527
528         Query can be a domain like 'bbc.co.uk'. Wildcards can be used, e.g. '*.bbc.co.uk'.
529         Alternatively, a query can contain a full domain name and some or all of a URL:
530         e.g. '*.wikipedia.org/wiki/*'
531
532         See <https://meta.wikimedia.org/wiki/Help:Linksearch> for details.
533
534         The generator returns dictionaries containing three keys:
535         - url: the URL linked to.
536         - ns: namespace of the wiki page
537         - pageid: the ID of the wiki page
538         - title: the page title.
539
540         """
541         self.require(1, 11)
542
543         kwargs = dict(listing.List.generate_kwargs('eu', query=query, prop=prop,
544                                                    protocol=protocol, namespace=namespace))
545         return listing.List(self, 'exturlusage', 'eu', limit=limit, **kwargs)
546
547     def logevents(self, type=None, prop=None, start=None, end=None,
548                   dir='older', user=None, title=None, limit=None, action=None):
549         self.require(1, 10)
550
551         kwargs = dict(listing.List.generate_kwargs('le', prop=prop, type=type, start=start,
552                                                    end=end, dir=dir, user=user, title=title, action=action))
553         return listing.List(self, 'logevents', 'le', limit=limit, **kwargs)
554
555     # def protectedtitles requires 1.15
556     def random(self, namespace, limit=20):
557         """Retrieves a generator of random page from a particular namespace.
558
559         limit specifies the number of random articles retrieved.
560         namespace is a namespace identifier integer.
561
562         Generator contains dictionary with namespace, page ID and title.
563
564         """
565         self.require(1, 12)
566
567         kwargs = dict(listing.List.generate_kwargs('rn', namespace=namespace))
568         return listing.List(self, 'random', 'rn', limit=limit, **kwargs)
569
570     def recentchanges(self, start=None, end=None, dir='older', namespace=None,
571                       prop=None, show=None, limit=None, type=None):
572         self.require(1, 9)
573
574         kwargs = dict(listing.List.generate_kwargs('rc', start=start, end=end, dir=dir,
575                                                    namespace=namespace, prop=prop, show=show, type=type))
576         return listing.List(self, 'recentchanges', 'rc', limit=limit, **kwargs)
577
578     def search(self, search, namespace='0', what='title', redirects=False, limit=None):
579         self.require(1, 11)
580
581         kwargs = dict(listing.List.generate_kwargs('sr', search=search, namespace=namespace, what=what))
582         if redirects:
583             kwargs['srredirects'] = '1'
584         return listing.List(self, 'search', 'sr', limit=limit, **kwargs)
585
586     def usercontributions(self, user, start=None, end=None, dir='older', namespace=None,
587                           prop=None, show=None, limit=None):
588         self.require(1, 9)
589
590         kwargs = dict(listing.List.generate_kwargs('uc', user=user, start=start, end=end,
591                                                    dir=dir, namespace=namespace, prop=prop, show=show))
592         return listing.List(self, 'usercontribs', 'uc', limit=limit, **kwargs)
593
594     def users(self, users, prop='blockinfo|groups|editcount'):
595         self.require(1, 12)
596
597         return listing.List(self, 'users', 'us', ususers='|'.join(users), usprop=prop)
598
599     def watchlist(self, allrev=False, start=None, end=None, namespace=None, dir='older',
600                   prop=None, show=None, limit=None):
601         self.require(1, 9)
602
603         kwargs = dict(listing.List.generate_kwargs('wl', start=start, end=end,
604                                                    namespace=namespace, dir=dir, prop=prop, show=show))
605         if allrev:
606             kwargs['wlallrev'] = '1'
607         return listing.List(self, 'watchlist', 'wl', limit=limit, **kwargs)
608
609     def expandtemplates(self, text, title=None, generatexml=False):
610         """Takes wikitext (text) and expands templates."""
611         self.require(1, 11)
612
613         kwargs = {}
614         if title is None:
615             kwargs['title'] = title
616         if generatexml:
617             kwargs['generatexml'] = '1'
618
619         result = self.api('expandtemplates', text=text, **kwargs)
620
621         if generatexml:
622             return result['expandtemplates']['*'], result['parsetree']['*']
623         else:
624             return result['expandtemplates']['*']
625
626     def ask(self, query, title=None):
627         """Ask a query against Semantic MediaWiki."""
628         kwargs = {}
629         if title is None:
630             kwargs['title'] = title
631         result = self.raw_api('ask', query=query, **kwargs)
632         return result['query']['results']

Benjamin Mako Hill || Want to submit a patch?