Pull links out of a Wikipedia XML dump, real fast.

wikipedia_links.py

#!/usr/bin/env python
# -*- coding: UTF-8 -*-

# Pull links out of a Wikipedia XML dump on STDIN and spit them out on
# STDOUT.


from mwlib.uparser import parseString as parseWikiMarkup
from xml.dom import pulldom
import sys
import time
import multiprocessing


process_count = multiprocessing.cpu_count() # xml parser competes with one queue
max_tasks = 1000
chunk_size = 100
page_notify = 1000

def find_links(xml):
    wm = parseWikiMarkup(title="", raw=xml)
    output = []
    for e in wm.allchildren():
        if e.type == e.t_http_url:
#            pprint.pprint(dict([(k,getattr(e,k)) for k in dir(e) ]))
            output.append(
                e.caption \
                .split('|', 1)[0] \
                .split('<', 1)[0] \
                .split('}', 1)[0] \
                .split("#", 1)[0] \
                .replace("&", "&") \
                .encode('utf-8')
            )
        elif e.type == e.t_complex_named_url:
            output.append(e.caption.strip().encode('utf-8'))
    return "\n".join(output)

def xml_parse(stream):
    doc = pulldom.parse(stream)
    i = 0
    s = time.time()
    for event, node in doc:
        if event == "START_ELEMENT":
            if node.tagName == "page":
                i += 1
                if i % page_notify == 0:
                    rate = page_notify / (time.time() - s)
                    s = time.time()
                    sys.stderr.write(
                        "%d pages parsed... (%i pages/sec)\n" % (i, rate)
                    )
            if node.tagName == "text":
                doc.expandNode(node)
                content = node.toxml().strip()
                node.unlink()
                if content:
                    yield content


pool = multiprocessing.Pool(
    processes=process_count, 
    maxtasksperchild=max_tasks
)
results = pool.imap_unordered(find_links, xml_parse(sys.stdin), chunk_size)

for r in results:
    if r:
        print r