random_pages.py

#!/usr/bin/env python

##########################################################################
#                                                                        #
# This program is free software; you can redistribute it and/or modify   #
# it under the terms of the GNU General Public License as published by   #
# the Free Software Foundation; version 2 of the License.                #
#                                                                        #
# This program is distributed in the hope that it will be useful,        #
# but WITHOUT ANY WARRANTY; without even the implied warranty of         #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the           #
# GNU General Public License for more details.                           #
#                                                                        #
##########################################################################

from sonet.mediawiki import HistoryPageProcessor, \
                            get_translations, get_tags, \
                            explode_dump_filename

import csv
import sys
import logging
from random import random
from datetime import datetime

from sonet import lib
from sonet.timr import Timr
from django.utils.encoding import smart_str


class HistoryRevisionsPageProcessor(HistoryPageProcessor):
    counter_saved = None
    queue = None
    output = None
    desired_page_type = None
    _counter_revisions = None
    _editors = None
    _min_text_length = None
    _initial_revision = None

    def __init__(self, **kwargs):
        super(HistoryRevisionsPageProcessor, self).__init__(**kwargs)
        self.counter_saved = 0
        self.queue = []

    def save(self):
        """
        it flushes the queue containing the processed and chosen pages
        on the csv output file
        """
        if len(self.queue):
            #logging.info('Saving %d pages' % len(self.queue))
            for q in self.queue:
                if q and q != '':
                    if self.output:
                        print >> self.output, q
                    else:
                        print q
            self.queue = []

    def end(self):
        logging.info('PAGES: %d - SAVED: %d' % (self.counter_pages,
                                                self.counter_saved))
        self.save()

    def process_title(self, elem):

        #clean attributes
        self.delattr(("_counter", "_type", "_title", "_editors",
                      "_initial_revision"))

        self._editors = []
        self._counter_revisions = 0
        self._min_text_length = 0

        a_title = elem.text.split(':')

        if len(a_title) == 1:
            if self.desired_page_type == "talk":
                self._skip = True
                return

            self._talk = False
            self._title = a_title[0]
        else:
            if a_title[0] == self.talkns and \
                    self.desired_page_type != "content":
                self._talk = True
                self._title = a_title[1]
            else:
                self._skip = True
                return

        self._desired = self.is_desired(self._title)

    def process_text(self, elem):
        if self._skip:
            return
        self._counter_revisions += 1
        if self.min_text and \
            len(elem.text) < self.min_text:
            if self._desired:
                logging.warning("Desired page %s skipped due to " \
                                "its text size: %d" %
                                (self._title, len(elem.text)))
            self._skip = True
            return
        else:
            pass

    def process_timestamp(self, elem):
        if self._skip:
            return

        if self.start_revision and not self._initial_revision:
            timestamp = elem.text
            year = int(timestamp[:4])
            month = int(timestamp[5:7])
            day = int(timestamp[8:10])
            revision_time = datetime(year, month, day)

            if revision_time > self.start_revision:
                if self._desired:
                    logging.warning("Desired page %s skipped due to " \
                                    "its initial revision date: %s" %
                                    (self._title,
                                     revision_time.strftime('%Y-%m-%d')))
                self._skip = True
                return

            self._initial_revision = revision_time

        else:
            pass

    def process_username(self, elem):
        if elem.text and elem.text != '':
            u = elem.text.encode('utf-8')
        else:
            return

        if u not in self._editors:
            self._editors.append(u)

    def process_ip(self, elem):
        if elem.text and elem.text != '':
            u = elem.text
        else:
            return

        if u not in self._editors:
            self._editors.append(u)

    def process_page(self, elem):

        if not self._skip and not self._desired and self.threshold < 1.:
            if random() > self.threshold:
                self._skip = True

        if not self._skip and self.n_users and \
           len(self._editors) < self.n_users:
            if self._desired:
                logging.warning("Desired page %s skipped due to low number "
                                "of editors: %d" %
                                (self._title, len(self._editors)))
            self._skip = True

        if not self._skip and self.min_revisions and \
           self._counter_revisions < self.min_revisions:
            if self._desired:
                logging.warning("Desired page %s skipped due to low number "
                                "of revisions: %d" %
                                (self._title, len(self._editors)))
            self._skip = True

        if not self._skip:
            self.queue.append(smart_str('%s%s' % (
                'Talk:' if self._talk else '', self._title,)))
            self.counter_saved += 1

        self.counter_pages += 1

        if not self.counter_pages % 10000:
            self.save()
            logging.info('PAGES: %d - SAVED: %d' %
                         (self.counter_pages, self.counter_saved))

        self._skip = False

    def process_redirect(self, _):
        # This class only considers pages that are in the desired file,
        # these pages must not be redirects
        self._skip = True
        #raise ValueError, "The page %s is a redirect. " % self._title + \
        #                  "Pages in the desired list must not be redirects."


def dumps_checker(args, type_):
    import re

    #logging.info(args)

    if args.editors_number or args.initial_revision or args.revisions_number:
        assert re.search('.-(meta-history)', type_), \
               "Wrong dump file, required: *-meta-history"

    if args.min_text_length:
        assert type_ == '-pages-meta-current', \
               "Wrong dump file, required: pages-meta-current"


def create_option_parser():
    import argparse

    p = argparse.ArgumentParser(description='Extract a random sample of pages '
                                            'title out of a Wikipedia dump')

    ## optional parameters
    p.add_argument('-t', '--type', default="content", metavar="TYPE",
                   help="Type of page to analize (content|talk|all)"
                        " - default: %(default)s")
    p.add_argument('-e', '--editors-number', default=0, metavar="NUM_EDITORS",
                   type=int, help="pages with less than NUM_EDITORS editors "
                                  "are skipped (default: %(default)s)")
    p.add_argument('-n', '--revisions-number', default=0,
                   metavar="NUM_REVISIONS", type=int,
                   help="pages with less than NUM_REVISIONS revisions "
                        "are skipped (default: %(default)s)")
    p.add_argument('-s', '--initial_revision', type=lib.yyyymmdd_to_datetime,
                   metavar="YYYYMMDD", help="Skip pages whose first revision "
                   "occurred later than this date", default=None)
    p.add_argument('-r', '--ratio', default=1., type=float, metavar="RATIO",
                  help="percentage of pages to be analyzed")
    p.add_argument('-T', '--min-text-length', default=0, metavar="TEXT_LENGTH",
                   type=int, help="pages with text shorter than TEXT_LENGTH "
                   "characters are skipped (default: %(default)s)")
    p.add_argument('-o', '--output', help="output file name",
                   metavar="OUTPUT_FILE", default=None)
    p.add_argument('--encoding', action="store", dest="encoding",
                 default="latin-1", help="encoding of the desired_list file")

    ## positional arguments
    p.add_argument('xml_fn', help="wikipedia dump to be parsed",
                   metavar="DUMP_FILE")
    p.add_argument('desired_pages_fn',
                   help="csv file containing desired pages",
                   metavar="DESIRED_PAGES_FILE")
    return p


def main():

    logging.basicConfig(#filename="random_page_extractor.log",
                                stream=sys.stderr,
                                level=logging.DEBUG)

    op = create_option_parser()
    args = op.parse_args()

    lang, date_, type_ = explode_dump_filename(args.xml_fn)
    deflate, _lineno = lib.find_open_for_this_file(args.xml_fn)

    dumps_checker(args, type_)

    logging.info('---------------------START---------------------')

    if _lineno:
        src = deflate(args.xml_fn, 51)
    else:
        src = deflate(args.xml_fn)

    translation = get_translations(src)
    tag = get_tags(src, tags='page,title,redirect,text,username,ip,timestamp')

    src.close()
    src = deflate(args.xml_fn)

    output = open(args.output, 'w') if args.output else None

    processor = HistoryRevisionsPageProcessor(
                    tag=tag,
                    lang=lang,
                    output=output,
                    threshold=args.ratio,
                    min_text=args.min_text_length,
                    min_revisions=args.revisions_number,
                    n_users=args.editors_number,
                    start_revision=args.initial_revision)

    processor.talkns = translation['Talk']
    processor.desired_page_type = args.type
    processor.set_desired_from_csv(desired_pages_fn, encoding=args.encoding)
    with Timr('processing'):
        processor.start(src)


if __name__ == "__main__":
    main()