diff --git a/sc2/bot_ai_internal.py b/sc2/bot_ai_internal.py index 056fa2a6..8bf8fa9f 100644 --- a/sc2/bot_ai_internal.py +++ b/sc2/bot_ai_internal.py @@ -722,6 +722,7 @@ def _prepare_units(self: BotAI) -> None: self.reactor_tags: set[int] = set() worker_types: set[UnitTypeId] = {UnitTypeId.DRONE, UnitTypeId.DRONEBURROWED, UnitTypeId.SCV, UnitTypeId.PROBE} + structure_types: frozenset[int] = self.game_data.structure_types index: int = 0 for unit in self.state.observation_raw.units: @@ -760,7 +761,7 @@ def _prepare_units(self: BotAI) -> None: elif alliance == 1: self.all_own_units.append(unit_obj) unit_id: UnitTypeId = unit_obj.type_id - if unit_obj.is_structure: + if unit_type in structure_types: self.structures.append(unit_obj) if unit_id in race_townhalls[self.race]: self.townhalls.append(unit_obj) @@ -790,7 +791,7 @@ def _prepare_units(self: BotAI) -> None: # Alliance.Enemy.value = 4 elif alliance == 4: self.all_enemy_units.append(unit_obj) - if unit_obj.is_structure: + if unit_type in structure_types: self.enemy_structures.append(unit_obj) else: self.enemy_units.append(unit_obj) diff --git a/sc2/game_data.py b/sc2/game_data.py index 1a590952..c3acef62 100644 --- a/sc2/game_data.py +++ b/sc2/game_data.py @@ -30,6 +30,10 @@ def __init__(self, data: sc2api_pb2.ResponseData) -> None: } self.units: dict[int, UnitTypeData] = {u.unit_id: UnitTypeData(self, u) for u in data.units if u.available} self.upgrades: dict[int, UpgradeData] = {u.upgrade_id: UpgradeData(self, u) for u in data.upgrades} + # Precomputed structure type IDs (raw ints) for O(1) membership checks + self.structure_types: frozenset[int] = frozenset( + u.unit_id for u in data.units if u.available and Attribute.Structure.value in u.attributes + ) # Cached UnitTypeIds so that conversion does not take long. This needs to be moved elsewhere if a new GameData object is created multiple times per game @lru_cache(maxsize=256) diff --git a/sc2/unit.py b/sc2/unit.py index 73c519b0..466db93c 100644 --- a/sc2/unit.py +++ b/sc2/unit.py @@ -34,7 +34,6 @@ IS_RETURNING, IS_REVEALED, IS_SNAPSHOT, - IS_STRUCTURE, IS_VISIBLE, OFF_CREEP_SPEED_INCREASE_DICT, OFF_CREEP_SPEED_UPGRADE_DICT, @@ -162,7 +161,7 @@ def tag(self) -> int: @property def is_structure(self) -> bool: """Checks if the unit is a structure.""" - return IS_STRUCTURE in self._type_data._proto.attributes + return self._proto.unit_type in self._bot_object.game_data.structure_types @property def is_light(self) -> bool: