diff --git a/sc2/constants.py b/sc2/constants.py index a3be87c3..cbbdd667 100644 --- a/sc2/constants.py +++ b/sc2/constants.py @@ -168,6 +168,7 @@ IS_PSIONIC: int = Attribute.Psionic.value UNIT_BATTLECRUISER: UnitTypeId = UnitTypeId.BATTLECRUISER UNIT_ORACLE: UnitTypeId = UnitTypeId.ORACLE +UNIT_VOIDRAY: UnitTypeId = UnitTypeId.VOIDRAY TARGET_GROUND: set[int] = {TargetType.Ground.value, TargetType.Any.value} TARGET_AIR: set[int] = {TargetType.Air.value, TargetType.Any.value} TARGET_BOTH: set[int] = TARGET_GROUND | TARGET_AIR diff --git a/sc2/unit.py b/sc2/unit.py index 73c519b0..5cd60ddf 100644 --- a/sc2/unit.py +++ b/sc2/unit.py @@ -50,6 +50,7 @@ UNIT_COLOSSUS, UNIT_ORACLE, UNIT_PHOTONCANNON, + UNIT_VOIDRAY, transforming, ) from sc2.data import Attribute, CloakState, Race, Target, race_gas, warpgate_abilities @@ -217,7 +218,7 @@ def _weapons(self): def can_attack(self) -> bool: """Checks if the unit can attack at all.""" # TODO BATTLECRUISER doesnt have weapons in proto?! - return bool(self._weapons) or self.type_id in {UNIT_BATTLECRUISER, UNIT_ORACLE} + return bool(self._weapons) or self.type_id in {UNIT_BATTLECRUISER, UNIT_ORACLE, UNIT_VOIDRAY} @property def can_attack_both(self) -> bool: @@ -227,7 +228,7 @@ def can_attack_both(self) -> bool: @cached_property def can_attack_ground(self) -> bool: """Checks if the unit can attack ground units.""" - if self.type_id in {UNIT_BATTLECRUISER, UNIT_ORACLE}: + if self.type_id in {UNIT_BATTLECRUISER, UNIT_ORACLE, UNIT_VOIDRAY}: return True if self._weapons: return any(weapon.type in TARGET_GROUND for weapon in self._weapons) @@ -247,7 +248,7 @@ def ground_range(self) -> float: """Returns the range against ground units. Does not include upgrades.""" if self.type_id == UNIT_ORACLE: return 4 - if self.type_id == UNIT_BATTLECRUISER: + if self.type_id in {UNIT_BATTLECRUISER, UNIT_VOIDRAY}: return 6 if self.can_attack_ground: weapon = next((weapon for weapon in self._weapons if weapon.type in TARGET_GROUND), None) @@ -258,7 +259,7 @@ def ground_range(self) -> float: @cached_property def can_attack_air(self) -> bool: """Checks if the unit can air attack at all. Does not include upgrades.""" - if self.type_id == UNIT_BATTLECRUISER: + if self.type_id in {UNIT_BATTLECRUISER, UNIT_VOIDRAY}: return True if self._weapons: return any(weapon.type in TARGET_AIR for weapon in self._weapons) @@ -276,7 +277,7 @@ def air_dps(self) -> float: @cached_property def air_range(self) -> float: """Returns the range against air units. Does not include upgrades.""" - if self.type_id == UNIT_BATTLECRUISER: + if self.type_id in {UNIT_BATTLECRUISER, UNIT_VOIDRAY}: return 6 if self.can_attack_air: weapon = next((weapon for weapon in self._weapons if weapon.type in TARGET_AIR), None) diff --git a/test/test_pickled_data.py b/test/test_pickled_data.py index fac77b71..196c57b4 100644 --- a/test/test_pickled_data.py +++ b/test/test_pickled_data.py @@ -782,6 +782,24 @@ def test_unit(): # assert marauder1.calculate_damage_vs_target(marauder_15_hp, include_overkill_damage=False)[0] == 15 +def test_units_without_weapon_data(): + """The game data lists no weapon for these units; their attacks are hardcoded.""" + bot: BotAI = get_map_specific_bot(random.choice(MAPS)) + proto = type(bot.workers.random._proto)() + proto.CopyFrom(bot.workers.random._proto) + proto.unit_type = UnitTypeId.VOIDRAY.value + void_ray = Unit(proto, bot) + try: + assert not void_ray._weapons + assert void_ray.can_attack + assert void_ray.can_attack_both + assert void_ray.ground_range == 6 + assert void_ray.air_range == 6 + finally: + # test_units counts the entries in this shared cache + Unit.class_cache.pop(UnitTypeId.VOIDRAY.value, None) + + def test_units(): bot: BotAI = get_map_specific_bot(random.choice(MAPS)) scvs = bot.workers