SC2: QoL for excluding all variants of a mission

worlds/sc2/mission_tables.py

@@ -297,6 +297,13 @@ class MissionInfo(NamedTuple):
         short_name = mission.mission_name[:mission.mission_name.find(' (')]
         lookup_name_to_mission[short_name] = mission
 
+lookup_short_name_to_mission_variants: Dict[str, Set[SC2Mission]] = {
+    mission.mission_name[:mission.mission_name.find(' (')]: {mission} for mission in SC2Mission if mission.flags & MissionFlag.HasRaceSwap
+}
+for mission in SC2Mission:
+    if mission.flags & MissionFlag.RaceSwap:
+        lookup_short_name_to_mission_variants[mission.mission_name[:mission.mission_name.find(' (')]].add(mission)
+
 lookup_id_to_campaign: Dict[int, SC2Campaign] = {
     campaign.id: campaign for campaign in SC2Campaign
 }

worlds/sc2/options.py

@@ -5,8 +5,9 @@
 from Options import *
 from Utils import get_fuzzy_results
 from BaseClasses import PlandoOptions
-from .mission_tables import SC2Campaign, SC2Mission, lookup_name_to_mission, MissionPools, get_missions_with_any_flags_in_list, \
-    campaign_mission_table, SC2Race, MissionFlag
+from .mission_tables import SC2Campaign, SC2Mission, lookup_name_to_mission, MissionPools, \
+    get_missions_with_any_flags_in_list, \
+    campaign_mission_table, SC2Race, MissionFlag, lookup_short_name_to_mission_variants
 from .mission_groups import mission_groups, MissionGroupNames
 from .mission_order.options import CustomMissionOrder
 from . import item_names
@@ -806,7 +807,7 @@ class ExcludedMissions(Sc2MissionSet):
     Doesn't apply to vanilla mission order.
     It may be impossible to build a valid campaign if too many missions are excluded."""
     display_name = "Excluded Missions"
-    valid_keys = {mission.mission_name for mission in SC2Mission}
+    valid_keys = lookup_name_to_mission.keys()
 
 
 class ExcludeVeryHardMissions(Choice):
@@ -1114,6 +1115,10 @@ def get_excluded_missions(world: 'SC2World') -> Set[SC2Mission]:
 
     excluded_missions: Set[SC2Mission] = set([lookup_name_to_mission[name] for name in excluded_mission_names])
 
+    for name in excluded_mission_names:
+        if name in lookup_short_name_to_mission_variants.keys():
+            excluded_missions.update(lookup_short_name_to_mission_variants[name])
+
     # Excluding Very Hard missions depending on options
     if world.options.exclude_very_hard_missions == ExcludeVeryHardMissions.option_true or (
         world.options.exclude_very_hard_missions == ExcludeVeryHardMissions.option_default and (

worlds/sc2/test/test_generation.py

@@ -529,3 +529,34 @@ def test_nco_and_wol_picks_correct_starting_mission(self):
         self.generate_world(world_options)
         self.assertEqual(get_first_mission(self.world, self.world.custom_mission_order), mission_tables.SC2Mission.LIBERATION_DAY)
 
+    def test_excluding_mission_short_name_excludes_all_variants_of_mission(self):
+        world_options = {
+            'excluded_missions': [
+                mission_tables.SC2Mission.ZERO_HOUR.mission_name.split(" (")[0]
+            ],
+            'mission_order': options.MissionOrder.option_grid,
+            'selected_races': options.SelectRaces.option_all,
+            'enable_race_swap': options.EnableRaceSwapVariants.option_shuffle_all,
+        }
+        self.generate_world(world_options)
+        missions = get_all_missions(self.world.custom_mission_order)
+        self.assertTrue(missions)
+        self.assertNotIn(mission_tables.SC2Mission.ZERO_HOUR, missions)
+        self.assertNotIn(mission_tables.SC2Mission.ZERO_HOUR_Z, missions)
+        self.assertNotIn(mission_tables.SC2Mission.ZERO_HOUR_P, missions)
+
+    def test_excluding_mission_variant_excludes_just_that_variant(self):
+        world_options = {
+            'excluded_missions': [
+                mission_tables.SC2Mission.ZERO_HOUR.mission_name
+            ],
+            'mission_order': options.MissionOrder.option_grid,
+            'selected_races': options.SelectRaces.option_all,
+            'enable_race_swap': options.EnableRaceSwapVariants.option_shuffle_all,
+        }
+        self.generate_world(world_options)
+        missions = get_all_missions(self.world.custom_mission_order)
+        self.assertTrue(missions)
+        self.assertNotIn(mission_tables.SC2Mission.ZERO_HOUR, missions)
+        self.assertIn(mission_tables.SC2Mission.ZERO_HOUR_Z, missions)
+        self.assertIn(mission_tables.SC2Mission.ZERO_HOUR_P, missions)