_filename_lock = threading.RLock()

def __init__(self, queue, outputdir, maxfilesize, prefix, compress, output_format):

threading.Thread.__init__(self)

self._queue = queue

self._maxfilesize = maxfilesize

self._prefix = prefix

self._compress = compress

self._outputdir = outputdir

self._output_format = output_format

if not os.path.exists(outputdir):

os.mkdir(outputdir)

self._outfile = None

@classmethod

def _get_file(cls, outputdir, compress=False):

with cls._filename_lock:

fpath = None

while not fpath or os.path.exists(fpath):

fname = ''.join([random.choice(string.letters) for _ in range(16)])

ext = ".raw" if not compress else ".raw.bz2"

fpath = os.path.join(outputdir, fname + ext)

if compress:

return bz2.BZ2File(fpath, 'w')

return open(fpath, 'w')

def _geturl(self, wiki_id):

return "%s?curid=%s" % (self._prefix, wiki_id)

def _write(self, wiki_id, wiki_title, wiki_text):

if not self._outfile and self._output_format != LEVELDB :

self._outfile = self._get_file(self._outputdir, self._compress)

colon = wiki_title.find(':')

if colon < 0 or wiki_title[:colon] in acceptedNamespaces:

print "[%s] [%s]" % (wiki_id.encode('utf-8'), wiki_title.encode('utf-8'))

#store to leveldb

if self._output_format == LEVELDB:

text = ' '.join(compact(clean(wiki_text))).strip()

DB.put(wiki_title.encode("utf-8"), text.encode("utf-8"))

return

url = self._geturl(wiki_id)

if self._output_format == TANL:

header = '<doc id="%s" url="%s" title="%s">%s\n' % (wiki_id, url, wiki_title, wiki_title)

body = ' '.join(compact(clean(wiki_text))).strip()

footer = "\n</doc>"

self._outfile.write(header.encode("utf-8"))

self._outfile.write(body.encode("utf-8"))

self._outfile.write(footer.encode("utf-8"))

elif self._output_format == JSON:

article = dict(id=wiki_id, url=url, title=wiki_title, text=wiki_text)

self._outfile.write(json.dumps(article, encoding='utf-8') + '\n')

if self._output_format != LEVELDB:

if self._outfile.tell() > self._maxfilesize:

self._outfile.close()

self._outfile = None

def _clean(self, page_elem):

# wiki xml dumps has namespace

# use xmlns from the page element

def TAG(tag):

return page_elem.tag.split("page")[0] + tag

wiki_id = page_elem.find(TAG("id")).text.strip()

wiki_title = page_elem.find(TAG("title")).text.strip()

revision_elem = page_elem.find(TAG("revision"))

if revision_elem is not None:

text_elem = revision_elem.find(TAG("text"))

if text_elem is not None:

wiki_text = ""

if text_elem.text is not None:

wiki_text = text_elem.text.strip()

self._write(wiki_id, wiki_title, wiki_text)

def run(self):

while True:

page_elem = None

try:

page_elem = self._queue.get(timeout=1)

if page_elem is not None:

self._clean(page_elem)

except Queue.Empty:

break

except:

traceback.print_exc(file=sys.stdout)

finally:

if page_elem is not None:

page_elem.clear()

self._queue.task_done()

'gallery', 'timeline', 'noinclude', 'pre',

'table', 'tr', 'td', 'th', 'caption',

'form', 'input', 'select', 'option', 'textarea',

'ul', 'li', 'ol', 'dl', 'dt', 'dd', 'menu', 'dir',

'ref', 'references', 'img', 'imagemap', 'source'

'a', 'b', 'big', 'blockquote', 'center', 'cite', 'div', 'em',

'font', 'h1', 'h2', 'h3', 'h4', 'hiero', 'i', 'kbd', 'nowiki',

'p', 'plaintext', 's', 'small', 'span', 'strike', 'strong',

'sub', 'sup', 'tt', 'u', 'var',

# remove leading whitespace and underscores

title = title.strip(' _')

# replace sequences of whitespace and underscore chars with a single space

title = re.compile(r'[\s_]+').sub(' ', title)

m = re.compile(r'([^:]*):(\s*)(\S(?:.*))').match(title)

if m:

prefix = m.group(1)

if m.group(2):

optionalWhitespace = ' '

else:

optionalWhitespace = ''

rest = m.group(3)

ns = prefix.capitalize()

if ns in acceptedNamespaces:

# If the prefix designates a known namespace, then it might be

# followed by optional whitespace that should be removed to get

# the canonical page name

# (e.g., "Category: Births" should become "Category:Births").

title = ns + ":" + rest.capitalize()

else:

# No namespace, just capitalize first letter.

# If the part before the colon is not a known namespace, then we must

# not remove the space after the colon (if any), e.g.,

# "3001: The_Final_Odyssey" != "3001:The_Final_Odyssey".

# However, to get the canonical page name we must contract multiple

# spaces into one, because

# "3001: The_Final_Odyssey" != "3001: The_Final_Odyssey".

title = prefix.capitalize() + ":" + optionalWhitespace + rest

else:

# no namespace, just capitalize first letter

title = title.capitalize();

def fixup(m):

text = m.group(0)

code = m.group(1)

try:

if text[1] == "#": # character reference

if text[2] == "x":

return unichr(int(code[1:], 16))

else:

return unichr(int(code))

else: # named entity

return unichr(name2codepoint[code])

except:

return text # leave as is

openRE = re.compile(openDelim)

closeRE = re.compile(closeDelim)

# partition text in separate blocks { } { }

matches = [] # pairs (s, e) for each partition

nest = 0 # nesting level

start = openRE.search(text, 0)

if not start:

return text

end = closeRE.search(text, start.end())

next = start

while end:

next = openRE.search(text, next.end())

if not next: # termination

while nest: # close all pending

nest -=1

end0 = closeRE.search(text, end.end())

if end0:

end = end0

else:

break

matches.append((start.start(), end.end()))

break

while end.end() < next.start():

# { } {

if nest:

nest -= 1

# try closing more

last = end.end()

end = closeRE.search(text, end.end())

if not end: # unbalanced

if matches:

span = (matches[0][0], last)

else:

span = (start.start(), last)

matches = [span]

break

else:

matches.append((start.start(), end.end()))

# advance start, find next close

start = next

end = closeRE.search(text, next.end())

break # { }

if next != start:

# { { }

nest += 1

# collect text outside partitions

res = ''

start = 0

for s, e in matches:

res += text[start:s]

start = e

res += text[start:]

"""Drop from text the blocks identified in matches"""

matches.sort()

res = ''

start = 0

for s, e in matches:

res += text[start:s]

start = e

res += text[start:]

global keepLinks

link = match.group(1)

colon = link.find(':')

if colon > 0 and link[:colon] not in acceptedNamespaces:

return ''

trail = match.group(3)

anchor = match.group(2)

if not anchor:

anchor = link

anchor += trail

if keepLinks:

return '<a href="%s">%s</a>' % (link, anchor)

else:

# FIXME: templates should be expanded

# Drop transclusions (template, parser functions)

# See: http://www.mediawiki.org/wiki/Help:Templates

text = dropNested(text, r'{{', r'}}')

# Drop tables

text = dropNested(text, r'{\|', r'\|}')

# Expand links

text = wikiLink.sub(make_anchor_tag, text)

# Drop all remaining ones

text = parametrizedLink.sub('', text)

# Handle external links

text = externalLink.sub(r'\1', text)

text = externalLinkNoAnchor.sub('', text)

# Handle bold/italic/quote

text = bold_italic.sub(r'\1', text)

text = bold.sub(r'\1', text)

text = italic_quote.sub(r'&quot;\1&quot;', text)

text = italic.sub(r'&quot;\1&quot;', text)

text = quote_quote.sub(r'\1', text)

text = text.replace("'''", '').replace("''", '&quot;')

################ Process HTML ###############

# turn into HTML

text = unescape(text)

# do it again (&amp;nbsp;)

text = unescape(text)

# Collect spans

matches = []

# Drop HTML comments

for m in comment.finditer(text):

matches.append((m.start(), m.end()))

# Drop self-closing tags

for pattern in selfClosing_tag_patterns:

for m in pattern.finditer(text):

matches.append((m.start(), m.end()))

# Drop ignored tags

for left, right in ignored_tag_patterns:

for m in left.finditer(text):

matches.append((m.start(), m.end()))

for m in right.finditer(text):

matches.append((m.start(), m.end()))

# Bulk remove all spans

text = dropSpans(matches, text)

# Cannot use dropSpan on these since they may be nested

# Drop discarded elements

for pattern in discard_element_patterns:

text = pattern.sub('', text)

# Expand placeholders

for pattern, placeholder in placeholder_tag_patterns:

index = 1

for match in pattern.finditer(text):

text = text.replace(match.group(), '%s_%d' % (placeholder, index))

index += 1

text = text.replace('<<', u'«').replace('>>', u'»')

#############################################

# Drop preformatted

# This can't be done before since it may remove tags

text = preformatted.sub('', text)

# Cleanup text

text = text.replace('\t', ' ')

text = spaces.sub(' ', text)

text = dots.sub('...', text)

text = re.sub(u' (,:\.\)\]»)', r'\1', text)

text = re.sub(u'(\[\(«) ', r'\1', text)

text = re.sub(r'\n\W+?\n', '\n', text) # lines with only punctuations

text = text.replace(',,', ',').replace(',.', '.')

"""Deal with headers, lists, empty sections, residuals of tables"""

page = [] # list of paragraph

headers = {} # Headers for unfilled sections

emptySection = False # empty sections are discarded

inList = False # whether opened <UL>

for line in text.split('\n'):

if not line:

continue

# Handle section titles

m = section.match(line)

if m:

title = m.group(2)

lev = len(m.group(1))

if keepSections:

page.append("<h%d>%s</h%d>" % (lev, title, lev))

if title and title[-1] not in '!?':

title += '.'

headers[lev] = title

# drop previous headers

for i in headers.keys():

if i > lev:

del headers[i]

emptySection = True

continue

# Handle page title

if line.startswith('++'):

title = line[2:-2]

if title:

if title[-1] not in '!?':

title += '.'

page.append(title)

# handle lists

elif line[0] in '*#:;':

if keepSections:

page.append("<li>%s</li>" % line[1:])

else:

continue

# Drop residuals of lists

elif line[0] in '{|' or line[-1] in '}':

continue

# Drop irrelevant lines

elif (line[0] == '(' and line[-1] == ')') or line.strip('.-') == '':

continue

elif len(headers):

items = headers.items()

items.sort()

for (i, v) in items:

page.append(v)

headers.clear()

page.append(line) # first line

emptySection = False

elif not emptySection:

page.append(line)

# we expects large dumps so we are using iterparse method

context = etree.iterparse(inputdump)

context = iter(context)

# discover prefix from the xml dump file

# /mediawiki/siteinfo/base

prefix = None

for event, elem in context:

if event == "end" and elem.tag.endswith("base"):

prefix = elem.text[:elem.text.rfind("/")]

break

print "base url: %s" % prefix

# initialize wiki page queue

queue = Queue.Queue(maxsize=1024)

# start worker threads

for _ in range(multiprocessing.cpu_count()):

cleaner = WikiCleanerThread(queue, outputdir, maxfilesize, prefix, compress, outformat)

cleaner.setDaemon(False)

cleaner.start()

# put element pages in the queue to be processed by the cleaner threads

for event, elem in context:

if event == "end" and elem.tag.endswith("page"):

queue.put(elem)

global keepLinks, keepSections

parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, description=__doc__)

parser.add_argument("wikidump", help="XML wiki dump file")

parser.add_argument("outputdir", help="output directory")

parser.add_argument("-w", "--overwrite", default=False, action="store_const", const=True, help="Overwrite existing output dir")

parser.add_argument("-b", "--bytes", default="25M", help="put specified bytes per output file (default is %(default)s)", metavar="n[KM]")

parser.add_argument("-c", "--compress", default=False, action="store_const", const=True, help="compress output files using bzip")

parser.add_argument("-l", "--links", default=False, action="store_const", const=True, help="preserve links")

parser.add_argument("-s", "--sections", default=False, action="store_const", const=True, help="preserve sections")

parser.add_argument("-f", "--format", choices=(TANL, JSON, LEVELDB), default=LEVELDB, help="choose output format default is %(default)s")

args = parser.parse_args()

keepLinks = args.links

keepSections = args.sections

# Minimum size of output files

min_file_size = 200 * 1024

try:

if args.bytes[-1] in 'kK':

file_size = int(args.bytes[:-1]) * 1024

elif args.bytes[-1] in 'mM':

file_size = int(args.bytes[:-1]) * 1024 * 1024

else:

file_size = int(args.bytes)

if file_size < min_file_size: raise ValueError()

except ValueError:

print >> sys.stderr, \

'Insufficient or invalid bytes size (minimum per output is %d bytes)' \

% min_file_size

return

if args.format.lower() == LEVELDB:

import leveldb

global DB, DB_MAX_BUFFER

DB = leveldb.DB(args.outputdir)

DB_MAX_BUFFER = 2048 # The max value of waiting flushing

else:

if not os.path.exists(args.outputdir):

os.makedirs(args.outputdir)

else:

if args.overwrite:

shutil.rmtree(args.outputdir)

os.makedirs(args.outputdir)

else:

raise ValueError("%s already exists, use --overwrite to recreate" % args.outputdir)

if args.wikidump.lower().endswith("bz2"):

with bz2.BZ2File(args.wikidump, 'r') as inputdump:

process_data(inputdump, args.outputdir, file_size, args.compress, args.format.lower())

else:

with open(args.wikidump, 'r') as inputdump: