Unit.has_buff and BotAI.get_available_abilities(unit)

examples/threebase_voidray.py

@@ -3,6 +3,7 @@
 import sc2
 from sc2 import Race, Difficulty
 from sc2.constants import *
+from sc2.ids.buff_id import BuffId
 from sc2.player import Bot, Computer
 
 class ThreebaseVoidrayBot(sc2.BotAI):
@@ -23,6 +24,15 @@ async def on_step(self, iteration):
         else:
             nexus = self.units(NEXUS).ready.random
 
+        if not nexus.has_buff(BuffId.CHRONOBOOSTENERGYCOST):
+            abilities = await self.get_available_abilities(nexus)
+            if AbilityId.CHRONOBOOSTENERGYCOST in abilities:
+                    await self.do(nexus(AbilityId.CHRONOBOOSTENERGYCOST, nexus))
+            else:
+                await self.chat_send("Can't cast chrono boost")
+        else:
+            await self.chat_send("Nexus is already boosted")
+
         for idle_worker in self.workers.idle:
             mf = self.state.mineral_field.closest_to(idle_worker)
             await self.do(idle_worker.gather(mf))

examples/zerg_rush.py

@@ -32,7 +32,8 @@ async def on_step(self, iteration):
             await self.do(zl.attack(target))
 
         for queen in self.units(QUEEN).idle:
-            if queen.energy >= 25: # Hard coded, since this is not (yet) available
+            abilities = await self.get_available_abilities(queen)
+            if AbilityId.INJECTLARVA in abilities:
                 await self.do(queen(INJECTLARVA, hatchery))
 
         if self.vespene >= 100:

sc2/bot_ai.py

@@ -73,6 +73,10 @@ def expansion_locations(self):
 
         return centers
 
+    async def get_available_abilities(self, unit):
+        # right know only checks cooldown, energy cost, and whether the ability has been researched
+        return await self._client.query_available_abilities(unit)
+
     async def expand_now(self, building=None, max_distance=10, location=None):
         if not building:
             building = self.townhalls.first.type_id

sc2/client.py

@@ -6,6 +6,9 @@
 )
 
 import logging
+
+from sc2.ids.ability_id import AbilityId
+
 logger = logging.getLogger(__name__)
 
 from .cache import method_cache_forever
@@ -160,6 +163,16 @@ async def query_building_placement(self, ability, positions, ignore_resources=Tr
         ))
         return [ActionResult(p.result) for p in result.query.placements]
 
+    async def query_available_abilities(self, unit):
+        assert isinstance(unit, Unit)
+        result = await self._execute(query=query_pb.RequestQuery(
+            abilities=[query_pb.RequestQueryAvailableAbilities(
+                unit_tag=unit.tag
+            )]
+        ))
+        return [AbilityId(a.ability_id) for a in result.query.abilities[0].abilities]
+
+
     async def chat_send(self, message, team_only):
         ch = ChatChannel.Team if team_only else ChatChannel.Broadcast
         r = await self._execute(action=sc_pb.RequestAction(

sc2/ids/ability_id.py

@@ -4,6 +4,8 @@
 import enum
 
 class AbilityId(enum.Enum):
+    CHRONOBOOSTENERGYCOST = 3755 # temporary > please see PR#24
+    NEXUSMASSRECALL = 3757 # temporary > please see PR#24
     INVALID = 0
     SMART = 1
     STOP_STOP = 4

sc2/ids/buff_id.py

@@ -4,6 +4,7 @@
 import enum
 
 class BuffId(enum.Enum):
+    CHRONOBOOSTENERGYCOST = 281 # temporary > please see PR#24
     INVALID = 0
     GRAVITONBEAM = 5
     GHOSTCLOAK = 6

sc2/unit.py

@@ -1,4 +1,5 @@
 from s2clientprotocol import sc2api_pb2 as sc_pb, raw_pb2 as raw_pb
+from sc2.ids.buff_id import BuffId
 
 from .position import Point3
 from .data import Alliance, Attribute, DisplayType
@@ -184,6 +185,11 @@ def train(self, unit, *args, **kwargs):
     def build(self, unit, *args, **kwargs):
         return self(self._game_data.units[unit.value].creation_ability.id, *args, **kwargs)
 
+    def has_buff(self, buff):
+        assert isinstance(buff, BuffId)
+
+        return buff.value in self._proto.buff_ids
+
     def attack(self, *args, **kwargs):
         return self(AbilityId.ATTACK, *args, **kwargs)