From bc4eb576440cf2a4b9b79150a9238878ce7cb3df Mon Sep 17 00:00:00 2001
From: EnvyDragon <138727357+EnvyDragon@users.noreply.github.com>
Date: Sun, 26 May 2024 00:14:48 -0400
Subject: [PATCH] sc2: race-filtering pr feedback

---
 worlds/sc2/mission_tables.py |  2 +-
 worlds/sc2/options.py        | 39 +++++++++++++-----------------------
 2 files changed, 15 insertions(+), 26 deletions(-)

diff --git a/worlds/sc2/mission_tables.py b/worlds/sc2/mission_tables.py
index 68cb01e5d815..30f502e14699 100644
--- a/worlds/sc2/mission_tables.py
+++ b/worlds/sc2/mission_tables.py
@@ -486,5 +486,5 @@ def get_campaign_potential_goal_missions(campaign: SC2Campaign) -> List[SC2Missi
     return missions
 
 
-def get_missions_with_flags(flags: MissionFlag) -> List[SC2Mission]:
+def get_missions_with_any_flags_in_list(flags: MissionFlag) -> List[SC2Mission]:
     return [mission for mission in SC2Mission if flags & mission.flags]
diff --git a/worlds/sc2/options.py b/worlds/sc2/options.py
index ec0fc2976ec2..f3db0b1749a3 100644
--- a/worlds/sc2/options.py
+++ b/worlds/sc2/options.py
@@ -6,7 +6,7 @@
     PerGameCommonOptions, Option, VerifyKeys)
 from Utils import get_fuzzy_results
 from BaseClasses import PlandoOptions
-from .mission_tables import SC2Campaign, SC2Mission, lookup_name_to_mission, MissionPools, get_missions_with_flags, \
+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_orders import vanilla_shuffle_order, mini_campaign_order
 from .mission_groups import mission_groups, MissionGroupNames
@@ -59,15 +59,14 @@ class SelectRaces(Choice):
     """
     Pick which factions' missions and items can be shuffled into the world.
     """
-    # bit 0: terran, bit 1: zerg, bit 2: protoss. all disabled means plando only
     display_name = "Select Playable Races"
-    option_all = 7
-    option_terran = 1
-    option_zerg = 2
-    option_protoss = 4
-    option_terran_and_zerg = 3
-    option_terran_and_protoss = 5
-    option_zerg_and_protoss = 6
+    option_all = (MissionFlag.Terran|MissionFlag.Zerg|MissionFlag.Protoss).value
+    option_terran = MissionFlag.Terran.value
+    option_zerg = MissionFlag.Zerg.value
+    option_protoss = MissionFlag.Protoss.value
+    option_terran_and_zerg = (MissionFlag.Terran|MissionFlag.Zerg).value
+    option_terran_and_protoss = (MissionFlag.Terran|MissionFlag.Protoss).value
+    option_zerg_and_protoss = (MissionFlag.Zerg|MissionFlag.Protoss).value
     option_plando = 0
     default = option_all
 
@@ -967,14 +966,11 @@ def get_enabled_races(world: 'SC2World') -> Set[SC2Race]:
     if selection == SelectRaces.option_all:
         return set(SC2Race)
     enabled = set()
-    # if bit 0, enable terran
-    if selection & 1:
+    if selection & MissionFlag.Terran:
         enabled.add(SC2Race.TERRAN)
-    # if bit 1, enable zerg
-    if selection & (1 << 1):
+    if selection & MissionFlag.Zerg:
         enabled.add(SC2Race.ZERG)
-    # if bit 2, enable protoss
-    if selection & (1 << 2):
+    if selection & MissionFlag.Protoss:
         enabled.add(SC2Race.PROTOSS)
     return enabled
 
@@ -1007,18 +1003,11 @@ def get_disabled_campaigns(world: 'SC2World') -> Set[SC2Campaign]:
 
 
 def get_disabled_flags(world: 'SC2World') -> MissionFlag:
-    excluded = 0
-    races = get_enabled_races(world)
-    # filter out missions based on disabled races
-    if SC2Race.TERRAN not in races:
-        excluded |= MissionFlag.Terran
-    if SC2Race.ZERG not in races:
-        excluded |= MissionFlag.Zerg
-    if SC2Race.PROTOSS not in races:
-        excluded |= MissionFlag.Protoss
+    excluded = (MissionFlag.Terran|MissionFlag.Zerg|MissionFlag.Protoss) ^ MissionFlag(get_option_value(world, "selected_races"))
     # filter out no-build missions
     if not get_option_value(world, "shuffle_no_build"):
         excluded |= MissionFlag.NoBuild
+    # TODO: add more flags to potentially exclude once we have a way to get that from the player
     return MissionFlag(excluded)
 
 
@@ -1043,7 +1032,7 @@ def get_excluded_missions(world: 'SC2World') -> Set[SC2Mission]:
         )
     # Omitting No-Build missions if not shuffling no-build
     if disabled_flags:
-        excluded_missions = excluded_missions.union(get_missions_with_flags(disabled_flags))
+        excluded_missions = excluded_missions.union(get_missions_with_any_flags_in_list(disabled_flags))
     # Omitting missions not in enabled campaigns
     for campaign in disabled_campaigns:
         excluded_missions = excluded_missions.union(campaign_mission_table[campaign])