Replaced ability generation script with full build data generation script

Author: frugs

generate_abilities.py generate_build_data.pygenerate_abilities.py renamed to 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,29 +60,38 @@ 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,
                         help='the path to the balance data export')
 
     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():