se-simulator/parsexml.py

import jsonlines
from lxml import etree

from utils import *


def parse_posts(inputdir, outputdir):
    i = 0
    skipped = 0
    iterator = etree.iterparse(inputdir + "/Posts.xml", events=("start", "end"))
    _, root = next(iterator)
    with jsonlines.open(outputdir + "/Questions.jsonl", mode="w") as questions, \
            jsonlines.open(outputdir + "/Answers.jsonl", mode="w") as answers, \
            jsonlines.open(outputdir + "/Titles.jsonl", "w") as titles:
        for event, element in iterator:
            title = element.get("Title")
            # if element.get("Score") and int(element.get("Score")) <= 10:
            #     skipped += 1
            #     element.clear()
            #     continue
            if title:
                titles.write(title)
            body = element.get("Body")
            if body:
                text = html2text(body)
                if element.get("PostTypeId") == "1":
                    questions.write(text)
                else:
                    answers.write(text)
            element.clear()
            if i % 100 == 0:
                root.clear()
                print(i, skipped, i + skipped, end="\r")
            i += 1
    print_stats(i, skipped)


def parse_comments(inputdir, outputdir):
    i = 0
    iterator = etree.iterparse(inputdir + "/Comments.xml", events=("start", "end"))
    _, root = next(iterator)
    with jsonlines.open(outputdir + "/Comments.jsonl", mode="w") as comments:
        for event, element in iterator:
            text = element.get("Text")
            if text:
                comments.write(text)
            element.clear()
            if i % 100 == 0:
                root.clear()
                print(i, end="\r")
            i += 1
    print_stats(i)


def parse_usernames(inputdir, outputdir):
    i = 0
    iterator = etree.iterparse(inputdir + "/Users.xml", events=("start", "end"))
    _, root = next(iterator)
    with jsonlines.open(outputdir + "/Usernames.jsonl", mode="w") as usernames:
        for event, element in iterator:
            displayname = element.get("DisplayName")
            if displayname:
                usernames.write(displayname)
            element.clear()
            if i % 100 == 0:
                root.clear()
                print(i, end="\r")
            i += 1
    print_stats(i)


if __name__ == "__main__":
    settings = get_settings(1)
    # parse_posts(settings)
    # parse_comments(settings)
    # parse_comments(settings)