extract_quizblocks.py

#!/bin/env python
ENGINE_FILE = 'Engine_Cytogenetics.xml'
QUIZBLOCKS_FILE = 'QuizBlocks_Cytogenetics.xml'

from lxml import etree


def extract_quizblocks():
    """
    extract quizblocks from unity engine.
    it fetches conversations, then follow the NextConvId.
    """

    tree = etree.parse(ENGINE_FILE)
    quizblocks = etree.Element('QuizBlocks')

    next_conv_id = None
    next_is_quiz = False
    has_conversation = False

    for element in tree.iter():
        if element.tag == 'Conversation':
            next_conv_id = element.attrib.get('NextConvId')
            next_is_quiz = next_conv_id and next_conv_id.startswith('Quiz')

            if next_is_quiz:
                has_conversation = True
                quiz_block_el = etree.SubElement(quizblocks, 'QuizBlock', **element.attrib)

        if has_conversation and element.tag == 'Quiz':
            if next_is_quiz:
                quiz_el = etree.SubElement(quiz_block_el, 'Quiz', **element.attrib)

            next_conv_id = element.attrib.get('NextConvId')
            next_is_quiz = next_conv_id and next_conv_id.startswith('Quiz')

        if has_conversation and element.tag == 'Options':
            options_el = etree.SubElement(quiz_el, 'Options', **element.attrib)

        if has_conversation and element.tag == 'Option':
            etree.SubElement(options_el, 'Option', **element.attrib)

    # remove last NextConvId
    for quiz in quizblocks.getchildren():
        try:
            print quiz.getchildren()[-1].attrib['NextConvId']
            del quiz.getchildren()[-1].attrib['NextConvId']
        except KeyError:
            pass

    with open(QUIZBLOCKS_FILE, 'w+') as f:
        f.write(etree.tostring(quizblocks, pretty_print=True))

if __name__ == "__main__":
    extract_quizblocks()