Build data generation script now merges new build data with existing build data correctly

It also now takes the sc2reader project dir as a parameter and writes build data files
directly to correct locations in the sc2reader project

Author: frugs

generate_build_data.py

@@ -155,47 +155,98 @@ def generate_build_data(balance_data_path):
     sorted_units = collections.OrderedDict(sorted(units.items(), key=lambda x: int(x[0])))
     sorted_abilities = collections.OrderedDict(sorted(abilities.items(), key=lambda x: int(x[0])))
 
-    unit_lookup = [[unit_name, unit_name] for _, unit_name in sorted_units.items()]
+    unit_lookup = dict((unit_name, unit_name) for _, unit_name in sorted_units.items())
 
     return sorted_units, sorted_abilities, unit_lookup, ability_lookup
 
 
+def combine_lookups(old_unit_lookup, old_ability_lookup, new_unit_lookup, new_ability_lookup):
+    unit_lookup = collections.OrderedDict(old_unit_lookup)
+    ability_lookup = collections.OrderedDict(old_ability_lookup)
+
+    # Just straightforwardly add any missing units
+    unit_lookup.update(new_unit_lookup)
+
+    # When merging old and new ability lookups, prefer to overwrite old cell data with new cell data when merging rows
+    # in the case of a key clash, but preserve old cell data if cell is empty in new ability lookup table
+    for ability_name, commands in new_ability_lookup.items():
+        if ability_name not in ability_lookup:
+            ability_lookup[ability_name] = commands
+        else:
+            for i, command in enumerate(commands):
+                # Pad potential gaps with empty commands
+                while len(ability_lookup[ability_name]) <= i:
+                    ability_lookup[ability_name].append("")
+
+                if command != "":
+                    ability_lookup[ability_name][i] = command
+
+    return unit_lookup, ability_lookup
+
+
 def main():
-    parser = argparse.ArgumentParser(description='Generate [BUILD_VERSION]_abilities.csv, [BUILD_VERSION]_units.csv,'
-                                                 ' ability_lookup.csv and unit_lookup.csv files from exported balance'
-                                                 ' data.')
+    parser = argparse.ArgumentParser(description='Generate and install new [BUILD_VERSION]_abilities.csv, '
+                                                 '[BUILD_VERSION]_units.csv, and update ability_lookup.csv and '
+                                                 'unit_lookup.csv files with any new units and ability commands.')
+    parser.add_argument('expansion', metavar='EXPANSION', type=str, choices=['WoL', 'HotS', 'LotV'],
+                        help='the expansion level of the balance data export, one of \'WoL\', \'HotS\', or \'LotV\'')
     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')
+    parser.add_argument('project_path', metavar='SC2READER_PROJECT_PATH', type=str,
+                        help='the path to the root of the sc2reader project directory')
 
     args = parser.parse_args()
 
-    units, abilities, unit_lookup, ability_lookup = generate_build_data(args.balance_data_path)
+    units, abilities, new_unit_lookup, new_ability_lookup = generate_build_data(args.balance_data_path)
 
     if not units or not abilities:
         parser.print_help()
         print("\n")
 
         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=',', lineterminator=os.linesep)
+    unit_lookup_path = os.path.join(args.project_path, 'sc2reader', 'data', 'unit_lookup.csv')
+    with open(unit_lookup_path, 'r') as file:
+        csv_reader = csv.reader(file, delimiter=',', lineterminator=os.linesep)
+        old_unit_lookup = collections.OrderedDict([(row[0], row[1]) for row in csv_reader if len(row) > 1])
+
+    ability_lookup_path = os.path.join(args.project_path, 'sc2reader', 'data', 'ability_lookup.csv')
+    with open(ability_lookup_path, 'r') as file:
+        csv_reader = csv.reader(file, delimiter=',', lineterminator=os.linesep)
+        old_ability_lookup = collections.OrderedDict([(row[0], row[1:]) for row in csv_reader if len(row) > 0])
+
+    if not old_unit_lookup or not old_ability_lookup:
+        parser.print_help()
+        print("\n")
+
+        raise ValueError("Could not find existing unit or ability lookups. Is the sc2reader project path correct?")
+
+    unit_lookup, ability_lookup = combine_lookups(
+        old_unit_lookup, old_ability_lookup, new_unit_lookup, new_ability_lookup)
+
+    units_file_path = os.path.join(
+        args.project_path, 'sc2reader', 'data', args.expansion, '{}_units.csv'.format(args.build_version))
+    with open(units_file_path, 'w') as file:
+        csv_writer = csv.writer(file, delimiter=',', lineterminator=os.linesep)
         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=',', lineterminator=os.linesep)
+    abilities_file_path = os.path.join(
+        args.project_path, 'sc2reader', 'data', args.expansion, '{}_abilities.csv'.format(args.build_version))
+    with open(abilities_file_path, 'w') as file:
+        csv_writer = csv.writer(file, delimiter=',', lineterminator=os.linesep)
         for ability_index, ability_name in abilities.items():
             csv_writer.writerow([ability_index, ability_name])
 
-    with open('unit_lookup.csv'.format(args.build_version), 'w', newline='') as csvfile:
-        csv_writer = csv.writer(csvfile, delimiter=',', lineterminator=os.linesep)
-        for lookup_entry in unit_lookup:
-            csv_writer.writerow(lookup_entry)
+    with open(unit_lookup_path, 'w') as file:
+        csv_writer = csv.writer(file, delimiter=',', lineterminator=os.linesep)
+        for entry in unit_lookup.items():
+            csv_writer.writerow(list(entry))
 
-    with open('ability_lookup.csv'.format(args.build_version), 'w', newline='') as csvfile:
-        csv_writer = csv.writer(csvfile, delimiter=',', lineterminator=os.linesep)
+    with open(ability_lookup_path, 'w') as file:
+        csv_writer = csv.writer(file, delimiter=',', lineterminator=os.linesep)
         for ability_name, commands in ability_lookup.items():
             csv_writer.writerow([ability_name] + commands)