Merge pull request Ziktofel#271 from EnvyDragon/take-your-pick

SC2: Adding "Pick Just One" option for race-swaps

worlds/sc2/options.py

@@ -270,13 +270,12 @@ class EnableRaceSwapVariants(Choice):
     designed for. NOTE: Cutscenes are always skipped on race-swapped mission variants.
 
     Disabled: Don't shuffle any non-vanilla map variants into the pool.
+    Pick One: Shuffle up to 1 valid version of each map into the pool, depending on other settings.
     Shuffle All: Each version of a map can appear in the same pool (so a map can appear up to 3 times as different races)
-    ("Pick Just One At Random" coming soon)
     """
     display_name = "Enable Race-Swapped Mission Variants"
     option_disabled = 0
-    # TODO: Implement pick-one logic
-    # option_pick_one = 1
+    option_pick_one = 1
     option_shuffle_all = 2
     default = option_disabled
 
@@ -1111,6 +1110,21 @@ def get_excluded_missions(world: 'SC2World') -> Set[SC2Mission]:
     # Omitting missions not in enabled campaigns
     for campaign in disabled_campaigns:
         excluded_missions = excluded_missions.union(campaign_mission_table[campaign])
+    # Omitting unwanted mission variants
+    if get_option_value(world, "enable_race_swap") == EnableRaceSwapVariants.option_pick_one:
+        swaps = [
+            mission for mission in SC2Mission
+            if mission not in excluded_missions
+            and mission.flags & (MissionFlag.HasRaceSwap|MissionFlag.RaceSwap)
+        ]
+        while len(swaps) > 0:
+            curr = swaps[0]
+            variants = [mission for mission in swaps if mission.map_file == curr.map_file]
+            variants.sort(key=lambda mission: mission.id)
+            swaps = [mission for mission in swaps if mission not in variants]
+            if len(variants) > 1:
+                variants.pop(world.random.randint(0, len(variants)-1))
+                excluded_missions = excluded_missions.union(variants)
 
     return excluded_missions
 

worlds/sc2/test/test_usecases.py

@@ -4,6 +4,7 @@
 
 from .test_base import Sc2SetupTestBase
 from .. import get_all_missions, item_groups, item_names, items, mission_tables, options
+from ..mission_tables import vanilla_mission_req_table, SC2Campaign, SC2Race
 
 
 class TestSupportedUseCases(Sc2SetupTestBase):
@@ -260,3 +261,24 @@ def test_excluding_faction_on_vanilla_order_excludes_epilogue(self) -> None:
             self.assertNotIn(mission_tables.lookup_name_to_mission[region].campaign,
                              ([mission_tables.SC2Campaign.EPILOGUE]),
                              f"{region} is an epilogue mission!")
+
+    def test_race_swap_pick_one_has_correct_length_and_includes_swaps(self) -> None:
+        world_options = {
+            'selected_races': options.SelectRaces.option_all,
+            'enable_race_swap': options.EnableRaceSwapVariants.option_pick_one,
+            'enable_wol_missions': True,
+            'enable_nco_missions': False,
+            'enable_prophecy_missions': False,
+            'enable_hots_missions': False,
+            'enable_lotv_prologue_missions': False,
+            'enable_lotv_missions': False,
+            'enable_epilogue_missions': False,
+            'mission_order': options.MissionOrder.option_grid,
+            'excluded_missions': [mission_tables.SC2Mission.ZERO_HOUR.mission_name],
+        }
+        self.generate_world(world_options)
+        world_regions = [region.name for region in self.multiworld.regions]
+        world_regions.remove('Menu')
+        self.assertEqual(len(world_regions), len(vanilla_mission_req_table.get(SC2Campaign.WOL)))
+        races = set(mission_tables.lookup_name_to_mission[mission].race for mission in world_regions)
+        self.assertTrue(SC2Race.ZERG in races or SC2Race.PROTOSS in races)