From 7a163541f34f5c2743df1044586957ec5f43318c Mon Sep 17 00:00:00 2001 From: SamOgon-one Date: Thu, 5 Nov 2026 08:52:28 +0100 Subject: [PATCH 1/2] perf: precompute structure type ids in GameData for fast O(1) lookups --- sc2/bot_ai_internal.py | 5 +++-- sc2/game_data.py | 4 ++++ sc2/unit.py | 2 +- 3 files changed, 8 insertions(+), 3 deletions(-) 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..0f4d8e12 100644 --- a/sc2/unit.py +++ b/sc2/unit.py @@ -162,7 +162,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: From 80c187393fbeb664730278fc25028b41691866ee Mon Sep 17 00:00:00 2001 From: SamOgon-one <114773312+SamOgon-one@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:21:44 +0100 Subject: [PATCH 2/2] Remove IS_STRUCTURE constant Removed IS_STRUCTURE constant from unit.py --- sc2/unit.py | 1 - 1 file changed, 1 deletion(-) diff --git a/sc2/unit.py b/sc2/unit.py index 0f4d8e12..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,