diff --git a/worlds/sc2/mission_tables.py b/worlds/sc2/mission_tables.py index 7e4e2ee48758..36ca49bca234 100644 --- a/worlds/sc2/mission_tables.py +++ b/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 } diff --git a/worlds/sc2/options.py b/worlds/sc2/options.py index 39204e8ee79b..dc4e58db1187 100644 --- a/worlds/sc2/options.py +++ b/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 ( diff --git a/worlds/sc2/test/test_generation.py b/worlds/sc2/test/test_generation.py index d13aa84ae2c5..bfda3ee8e143 100644 --- a/worlds/sc2/test/test_generation.py +++ b/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) \ No newline at end of file