diff --git a/generate_abilities.py b/generate_build_data.py similarity index 73% rename from generate_abilities.py rename to generate_build_data.py index 83356e5d..3f039b38 100644 --- a/generate_abilities.py +++ b/generate_build_data.py @@ -6,8 +6,9 @@ import xml.etree.ElementTree -def generate_abilities(balance_data_path): +def generate_build_data(balance_data_path): abilities = {} + units = {} for xml_file_path in glob.glob(os.path.join(balance_data_path, "*.xml")): tree = xml.etree.ElementTree.parse(xml_file_path) @@ -19,8 +20,13 @@ def generate_abilities(balance_data_path): unit_name = root.get("id") + meta_element = root.find("./meta") + + if unit_name and meta_element is not None and meta_element.get("index"): + units[meta_element.get("index")] = unit_name + build_unit_element = root.find("./builds/unit") - if build_unit_element: + if build_unit_element is not None: build_ability_index = build_unit_element.get("ability") if unit_name == "SCV": @@ -54,17 +60,21 @@ def generate_abilities(balance_data_path): abilities[train_ability_index] = "{}Train{}".format(unit_name, trained_unit_name) research_upgrade_element = root.find("./researches/upgrade") - if research_upgrade_element: + if research_upgrade_element is not None: research_ability_index = research_upgrade_element.get("ability") research_ability_name = "{}Research".format(unit_name) abilities[research_ability_index] = research_ability_name - return collections.OrderedDict(sorted(abilities.items(), key=lambda x: int(x[0]))) + sorted_abilities = collections.OrderedDict(sorted(abilities.items(), key=lambda x: int(x[0]))) + sorted_units = collections.OrderedDict(sorted(units.items(), key=lambda x: int(x[0]))) + + return sorted_units, sorted_abilities def main(): - parser = argparse.ArgumentParser(description='Generate a [BUILD]_abilities.csv file from exported balance data.') + parser = argparse.ArgumentParser(description='Generate [BUILD_VERSION]_abilities.csv and [BUILD_VERSION]_units.csv' + ' files from exported balance data.') parser.add_argument('build_version', metavar='BUILD_VERSION', type=int, help='the build version of the balance data export') parser.add_argument('balance_data_path', metavar='BALANCE_DATA_PATH', type=str, @@ -72,11 +82,16 @@ def main(): args = parser.parse_args() - abilities = generate_abilities(args.balance_data_path) + units, abilities = generate_build_data(args.balance_data_path) - if not abilities: + if not units or not abilities: raise ValueError("No balance data found at provided balance data path.") + with open('{}_units.csv'.format(args.build_version), 'w', newline='') as csvfile: + csv_writer = csv.writer(csvfile, delimiter=',') + for unit_index, unit_name in units.items(): + csv_writer.writerow([unit_index, unit_name]) + with open('{}_abilities.csv'.format(args.build_version), 'w', newline='') as csvfile: csv_writer = csv.writer(csvfile, delimiter=',') for ability_index, ability_name in abilities.items():