<commit_before>#!/usr/bin/env python
"""
Index or split Wikipedia XML Dump into alphabetical (gz) files

This expects to operate on the (currently ~12GB)
latest/enwiki-latest-pages-articles.xml.bz2
The process is interruptable and restartable.

INPUT
   fname  XML Dump filename
          see -h for options

OUTPUT
   index or split (gzip) files to dest/[0-9], dest/[A-Z]

References
    https://en.wikipedia.org/wiki/Wikipedia:Database_download
    https://en.wikipedia.org/wiki/Wikipedia:Vital_articles/Expanded
    http://effbot.org/zone/celementtree.htm
"""

from __future__ import print_function

__author__ = "siznax"
__version__ = "1 Oct 2015"

import argparse
import bz2
import gzip
import os
import re
import sys
import string
import time
import traceback

DEFAULT_CHUNK_KB = 1
ONE_KB = 1000
ONE_MB = 1000**2
MAX_MEGABYTES = 10

from wp_parser import WPLineParser


class IndexParser(WPLineParser):

    def __init__(self, dest, offset, split):
        WPLineParser.__init__(self)
        self._files = dict()
        self._paths = dict()
        self.bytes_read = 0
        self.dest = dest
        self.first_title = ""
        self.first_title_start = 0
        self.offset = offset
        self.split = split
        self.tell = 0
        self.title = ""
        self.title_start = 0
        self.titles = None
        self.titles_found = []
        self.titles_written = 0

    def add_file(self, key):
        path = "%s/%s" % (self.dest, key)
        print("+ open %s" % path)
        self._files[key] = file(path, 'w')
        self._paths[key] = path

    def open_files(self):
        if not self.dest:
            return
        os.mkdir(self.dest)
        for char in string.digits + string.ascii_uppercase:
            path = "%s/%s" % (self.dest, char)
            print("+ open %s" % path)
            self._files[char] = gzip.open(path, 'wb')
            self._paths[char] = path
        if self.titles:
            self.add_file("titles_found")
            self.add_file("titles_left")

    def close_files(self):
        if not self.dest:
            return
        for key, handle in self._files.iteritems():
            tell = handle.tell()
            path = self._paths[key]
            handle.close()
            if tell:
                print("wrote %d bytes to %s" % (tell, path))

    def _write(self, title, title_start, elem):
        first_char = ascii_bin(title)
        if first_char in self._files:
            _file = self._files[first_char]
            if self.split:
                if self.titles:
                    if title in self.titles:
                        _file.write(elem)
                        self.titles_written += 1
                        self.titles_found.append(title)
                        self.titles.remove(title)
                else:
                    _file.write(elem)
            else:
                _file.write("%s %s\n" % (title, title_start))
        else:
            print("ORPHAN %s" % title)

    def process(self, elem):
        title = page_title(elem)
        title_start = self.offset + self.byte_count - len(elem)
        if not self.title:
            self.first_title = title
            self.first_title_start = title_start
        self.title = title
        self.title_start = title_start
        if not self.dest:
            # print("%s %s" % (title, len(elem)))
            return
        self._write(title, title_start, elem)


def page_title(page):
    """returns title from XML Dump <page>"""
    m = re.search(r"<title>([^<]*)<", page)
    if not m:
        print(page[:64])
        raise ValueError("title not found!")
    return m.group(1)


def ascii_bin(title):
    """returns first 0-9 or A-Z character in upper(title)"""
    try:
        return [x for x in title.upper()
                if 47 < ord(x) < 58 or 64 < ord(x) < 91][0]
    except:
        return title


def setup(dest, offset, split, titles):
    if dest and os.path.exists(dest):
        print("Destination exists: %s" % dest, file=sys.stderr)
        sys.exit(os.EX_IOERR)
    ip = IndexParser(dest, offset, split)
    ip.open_files()
    if titles:
        with open(titles) as fh:
            ip.titles = set(fh.read().split("\n"))
            print("Pulling %d titles" % len(ip.titles), file=sys.stderr)
    return ip


def gobble(ip, fname, chunk_size, max_mb, offset):
    chunk_size = ONE_KB * chunk_size
    est_bytes_read = 0
    max_bytes = ONE_MB * max_mb
    with bz2.BZ2File(fname, 'r') as zh:
        zh.seek(offset)
        try:
            while ip.bytes_read < max_bytes:
                # data = zh.read(chunk_size)
                data = zh.readlines(chunk_size)
                ip.parse(data)
                ip.tell = zh.tell()
                ip.bytes_read = ip.tell - offset
                est_bytes_read += chunk_size
                if est_bytes_read % ONE_MB * 10 == 0:
                    print("  %s %d" % (ip.title, ip.title_start))
        except KeyboardInterrupt:
            teardown(ip)<commit_msg>fix titles open_files(), flush stdout<commit_after>#!/usr/bin/env python
"""
Index or split Wikipedia XML Dump into alphabetical (gz) files

This expects to operate on the (currently ~12GB)
latest/enwiki-latest-pages-articles.xml.bz2
The process is interruptable and restartable.

INPUT
   fname  XML Dump filename
          see -h for options

OUTPUT
   index or split (gzip) files to dest/[0-9], dest/[A-Z]

References
    https://en.wikipedia.org/wiki/Wikipedia:Database_download
    https://en.wikipedia.org/wiki/Wikipedia:Vital_articles/Expanded
    http://effbot.org/zone/celementtree.htm
"""

from __future__ import print_function

__author__ = "siznax"
__version__ = "1 Oct 2015"

import argparse
import bz2
import gzip
import os
import re
import sys
import string
import time
import traceback

DEFAULT_CHUNK_KB = 1
ONE_KB = 1000
ONE_MB = 1000**2
MAX_MEGABYTES = 10

from wp_parser import WPLineParser


class IndexParser(WPLineParser):

    def __init__(self, dest, offset, split):
        WPLineParser.__init__(self)
        self._files = dict()
        self._paths = dict()
        self.bytes_read = 0
        self.dest = dest
        self.first_title = ""
        self.first_title_start = 0
        self.offset = offset
        self.split = split
        self.tell = 0
        self.title = ""
        self.title_start = 0
        self.titles = None
        self.titles_found = []
        self.titles_written = 0

    def add_file(self, key):
        path = "%s/%s" % (self.dest, key)
        print("+ open %s" % path)
        self._files[key] = file(path, 'w')
        self._paths[key] = path

    def open_files(self, titles=False):
        if not self.dest:
            return
        os.mkdir(self.dest)
        for char in string.digits + string.ascii_uppercase:
            path = "%s/%s" % (self.dest, char)
            print("+ open %s" % path)
            self._files[char] = gzip.open(path, 'wb')
            self._paths[char] = path
        if titles:
            self.add_file("titles_found")
            self.add_file("titles_left")

    def close_files(self):
        if not self.dest:
            return
        for key, handle in self._files.iteritems():
            tell = handle.tell()
            path = self._paths[key]
            handle.close()
            if tell:
                print("wrote %d bytes to %s" % (tell, path))

    def _write(self, title, title_start, elem):
        first_char = ascii_bin(title)
        if first_char in self._files:
            _file = self._files[first_char]
            if self.split:
                if self.titles:
                    if title in self.titles:
                        _file.write(elem)
                        self.titles_written += 1
                        self.titles_found.append(title)
                        self.titles.remove(title)
                else:
                    _file.write(elem)
            else:
                _file.write("%s %s\n" % (title, title_start))
        else:
            print("ORPHAN %s" % title)

    def process(self, elem):
        title = page_title(elem)
        title_start = self.offset + self.byte_count - len(elem)
        if not self.title:
            self.first_title = title
            self.first_title_start = title_start
        self.title = title
        self.title_start = title_start
        if not self.dest:
            # print("%s %s" % (title, len(elem)))
            return
        self._write(title, title_start, elem)


def page_title(page):
    """returns title from XML Dump <page>"""
    m = re.search(r"<title>([^<]*)<", page)
    if not m:
        print(page[:64])
        raise ValueError("title not found!")
    return m.group(1)


def ascii_bin(title):
    """returns first 0-9 or A-Z character in upper(title)"""
    try:
        return [x for x in title.upper()
                if 47 < ord(x) < 58 or 64 < ord(x) < 91][0]
    except:
        return title


def setup(dest, offset, split, titles):
    if dest and os.path.exists(dest):
        print("Destination exists: %s" % dest, file=sys.stderr)
        sys.exit(os.EX_IOERR)
    ip = IndexParser(dest, offset, split)
    ip.open_files(titles)
    if titles:
        with open(titles) as fh:
            ip.titles = set(fh.read().split("\n"))
            print("Pulling %d titles" % len(ip.titles), file=sys.stderr)
    return ip


def gobble(ip, fname, chunk_size, max_mb, offset):
    chunk_size = ONE_KB * chunk_size
    est_bytes_read = 0
    max_bytes = ONE_MB * max_mb
    with bz2.BZ2File(fname, 'r') as zh:
        zh.seek(offset)
        try:
            while ip.bytes_read < max_bytes:
                # data = zh.read(chunk_size)
                data = zh.readlines(chunk_size)
                ip.parse(data)
                ip.tell = zh.tell()
                ip.bytes_read = ip.tell - offset
                est_bytes_read += chunk_size
                if est_bytes_read % ONE_MB * 10 == 0:
                    print("  %s %d" % (ip.title, ip.title_start))
                    sys.stdout.flush()
        except KeyboardInterrupt:
            teardown(ip)
{'fraction_non_alphanumeric': 0.11082474226804123, 'fraction_numerical': 0.00971451229183188, 'mean_word_length': 1.8809251856082239, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 7, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 2}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2912453', 'n_tokens_mistral': 3324, 'n_tokens_neox': 3135, 'n_words': 791}