Python version of the `PD_extras_convert.bat` script originally created by "kdin6tl3nm"

pd_extras_convert.py

#!/usr/bin/env python3
import os
import sys
import getopt
import shutil
import zipfile
from lxml import etree as et


# Python version of the `PD_extras_convert.bat` script originally created by "kdin6tl3nm"
# Python version created by "methuselah"
# http://forum.pleasuredome.org.uk/index.php?showtopic=31185
#
# Usage
# =====
# pd_extras_convert.py -i <inputfile> -o <outputfile>
# pd_extras_convert.py -h
#
#
# Dependencies
# ============
# - python3
# - python3-lxml
def main(argv):
    in_file, out_file = parse_args(argv)
    tmp_dir = os.path.join(os.getcwd(), r"tmp")
    unzip_dats(tmp_dir, in_file)
    convert_extras_all_nzipped_content(tmp_dir, out_file)
    cleanup(tmp_dir)


def parse_args(argv):
    in_file = ""
    out_file = ""
    try:
        opts, args = getopt.getopt(argv, "hi:o:", ["ifile=", "ofile="])
    except getopt.GetoptError:
        print("pd_extras_convert.py -i <inputfile> -o <outputfile>")
        sys.exit(2)
    for opt, arg in opts:
        if opt == "-h":
            print("pd_extras_convert.py -i <inputfile> -o <outputfile>")
            sys.exit()
        elif opt in ("-i", "--ifile"):
            in_file = arg
        elif opt in ("-o", "--ofile"):
            out_file = arg

    if not in_file:
        print("input file required")
        sys.exit(2)

    if not out_file:
        print("output file required")
        sys.exit(2)

    return in_file, out_file


def unzip_dats(tmp_dir, dats_zip_file):
    # Clean out old tmp folder
    if os.path.exists(tmp_dir):
        shutil.rmtree(tmp_dir)

    # Create tmp folder
    os.makedirs(tmp_dir)

    # Extract
    with zipfile.ZipFile(dats_zip_file, "r") as zip_ref:
        zip_ref.extractall(tmp_dir)


def cleanup(tmp_dir):
    # Remove tmp folder
    shutil.rmtree(tmp_dir)


def convert_extras_samples(tmp_dir, parser):

    # Open original file
    xml_doc = et.parse(os.path.join(tmp_dir, r"samples.dat"), parser)
    samples_root = xml_doc.getroot()

    # Remove header
    for element in samples_root.iter("header"):
        samples_root.remove(element)

    # rename datafile to dir
    for element in samples_root.iter("datafile"):
        element.tag = "dir"
        element.set("name", "samples")

    return samples_root


def convert_extras_artwork(tmp_dir, parser):

    # Open original file
    xml_doc = et.parse(os.path.join(tmp_dir, r"artwork.dat"), parser)
    artwork_root = xml_doc.getroot()

    # Remove header
    for element in artwork_root.iter("header"):
        artwork_root.remove(element)

    # rename datafile to dir
    for element in artwork_root.iter("datafile"):
        element.tag = "dir"
        element.set("name", "artwork")

    return artwork_root


def wrap_up_machine_tag(root, match):
    match_list = root.xpath("//machine[@name='" + match + "']")
    if match_list:
        match = match_list[0]
        match.tag = "dir"
        new_tag = et.Element("game")
        new_tag.set("name", match.attrib["name"])
        new_tag.extend(match)
        match.append(new_tag)


def convert_extras_all_nzipped_content(tmp_dir, out_file):

    # Setup parser
    parser = et.XMLParser(remove_blank_text=True)

    # Open original file
    xml_doc = et.parse(os.path.join(tmp_dir, r"all_non-zipped_content.dat"), parser)
    root = xml_doc.getroot()

    # Remove existing header
    for element in root.iter("header"):
        root.remove(element)

    # Add new header
    header = et.Element("header")
    header_1 = et.SubElement(header, "name")
    header_1.text = "Extras"
    header_1 = et.SubElement(header, "description")
    header_1.text = "MAME Extras (all content)"
    header_1 = et.SubElement(header, "category")
    header_1.text = "Standard DatFile"
    et.SubElement(header, "version")
    header_1 = et.SubElement(header, "author")
    header_1.text = "newuzer"
    header_1 = et.SubElement(header, "comment")
    header_1.text = "For use with ROMVault"
    root.insert(0, header)

    # Wrap up ctrlr, dats, folders in their own dir tags
    # e.g. <dir name="ctrlr"><machine name="ctrlr">
    wrap_up_machine_tag(root, "ctrlr")
    wrap_up_machine_tag(root, "dats")
    wrap_up_machine_tag(root, "folders")
    wrap_up_machine_tag(root, "history")

    # Add artwork
    root.append(convert_extras_artwork(tmp_dir, parser))

    # Add samples
    root.append(convert_extras_samples(tmp_dir, parser))

    # rename machine to game
    for element in root.iter("machine"):
        element.tag = "game"

    # Write back to file
    xml_doc.write(out_file, xml_declaration=True, encoding="UTF-8", pretty_print=True)


if __name__ == "__main__":
    main(sys.argv[1:])