Add is_transforming and can_attack_both

sc2/constants.py

@@ -27,3 +27,69 @@
     PURIFIERVESPENEGEYSER.value,
     SHAKURASVESPENEGEYSER.value,
 }
+transforming = {
+    # terran structures
+    BARRACKS: LAND_BARRACKS,
+    BARRACKSFLYING: LAND_BARRACKS,
+    COMMANDCENTER: LAND_COMMANDCENTER,
+    COMMANDCENTERFLYING: LAND_COMMANDCENTER,
+    ORBITALCOMMAND: LAND_ORBITALCOMMAND,
+    ORBITALCOMMANDFLYING: LAND_ORBITALCOMMAND,
+    FACTORY: LAND_FACTORY,
+    FACTORYFLYING: LAND_FACTORY,
+    STARPORT: LAND_STARPORT,
+    STARPORTFLYING: LAND_STARPORT,
+    SUPPLYDEPOT: MORPH_SUPPLYDEPOT_RAISE,
+    SUPPLYDEPOTLOWERED: MORPH_SUPPLYDEPOT_LOWER,
+    # terran units
+    HELLION: MORPH_HELLION,
+    HELLIONTANK: MORPH_HELLBAT,
+    LIBERATOR: MORPH_LIBERATORAAMODE,
+    LIBERATORAG: MORPH_LIBERATORAGMODE,
+    SIEGETANK: UNSIEGE_UNSIEGE,
+    SIEGETANKSIEGED: SIEGEMODE_SIEGEMODE,
+    THOR: MORPH_THOREXPLOSIVEMODE,
+    THORAP: MORPH_THORHIGHIMPACTMODE,
+    VIKINGASSAULT: MORPH_VIKINGASSAULTMODE,
+    VIKINGFIGHTER: MORPH_VIKINGFIGHTERMODE,
+    WIDOWMINE: BURROWUP,
+    WIDOWMINEBURROWED: BURROWDOWN,
+    # protoss structures
+    GATEWAY: MORPH_GATEWAY,
+    WARPGATE: MORPH_WARPGATE,
+    # protoss units
+    OBSERVER: MORPH_OBSERVERMODE,
+    OBSERVERSIEGEMODE: MORPH_SURVEILLANCEMODE,
+    WARPPRISM: MORPH_WARPPRISMTRANSPORTMODE,
+    WARPPRISMPHASING: MORPH_WARPPRISMPHASINGMODE,
+    # zerg structures
+    SPINECRAWLER: SPINECRAWLERROOT_SPINECRAWLERROOT,
+    SPINECRAWLERUPROOTED: SPINECRAWLERUPROOT_SPINECRAWLERUPROOT,
+    SPORECRAWLER: SPORECRAWLERROOT_SPORECRAWLERROOT, 
+    SPORECRAWLERUPROOTED: SPORECRAWLERUPROOT_SPORECRAWLERUPROOT,
+    # zerg units
+    BANELING: BURROWUP_BANELING,
+    BANELINGBURROWED: BURROWDOWN_BANELING,
+    DRONE: BURROWUP_DRONE,
+    DRONEBURROWED: BURROWDOWN_DRONE,
+    HYDRALISK: BURROWUP_HYDRALISK,
+    HYDRALISKBURROWED: BURROWDOWN_HYDRALISK,
+    INFESTOR: BURROWUP_INFESTOR,
+    INFESTORBURROWED: BURROWDOWN_INFESTOR,
+    INFESTORTERRAN: BURROWUP_INFESTORTERRAN,
+    INFESTORTERRANBURROWED: BURROWDOWN_INFESTORTERRAN,
+    LURKERMP: BURROWUP_LURKER,
+    LURKERMPBURROWED: BURROWDOWN_LURKER,
+    OVERSEER: MORPH_OVERSEERMODE,
+    OVERSEERSIEGEMODE: MORPH_OVERSIGHTMODE,
+    QUEEN: BURROWUP_QUEEN,
+    QUEENBURROWED: BURROWDOWN_QUEEN,
+    ROACH: BURROWUP_ROACH,
+    ROACHBURROWED: BURROWDOWN_ROACH,
+    SWARMHOSTBURROWEDMP: BURROWDOWN_SWARMHOST,
+    SWARMHOSTMP: BURROWUP_SWARMHOST,
+    ULTRALISK: BURROWUP_ULTRALISK,
+    ULTRALISKBURROWED: BURROWDOWN_ULTRALISK,
+    ZERGLING: BURROWUP_ZERGLING,
+    ZERGLINGBURROWED: BURROWDOWN_ZERGLING,
+}

sc2/unit.py

@@ -3,6 +3,7 @@
 
 from . import unit_command
 from .cache import property_immutable_cache, property_mutable_cache
+from .constants import transforming
 from .data import Alliance, Attribute, CloakState, DisplayType, Race, TargetType, warpgate_abilities
 from .ids.ability_id import AbilityId
 from .ids.buff_id import BuffId
@@ -145,6 +146,11 @@ def can_attack(self) -> bool:
         # TODO BATTLECRUISER doesnt have weapons in proto?!
         return bool(self._weapons) or self.type_id == UnitTypeId.BATTLECRUISER
 
+    @property_immutable_cache
+    def can_attack_both(self) -> bool:
+        """ Checks if the unit can attack both ground and air units. """
+        return self.can_attack_ground and self.can_attack_air
+
     @property_immutable_cache
     def can_attack_ground(self) -> bool:
         """ Checks if the unit can attack ground units. """
@@ -584,6 +590,11 @@ def is_constructing_scv(self) -> bool:
             }
         )
 
+    @property_immutable_cache
+    def is_transforming(self) -> bool:
+        """ Checks if the unit transforming. """
+        return self.type_id in transforming and self.is_using_ability(transforming[self.type_id])
+
     @property_immutable_cache
     def is_repairing(self) -> bool:
         """ Checks if the unit is an SCV or MULE that is currently repairing. """