Warpgate support (#28)

Implement warpgate support
Fixes #27

examples/warpgate_toss.py

@@ -0,0 +1,115 @@
+import random
+
+import sc2
+from sc2 import Race, Difficulty
+from sc2.constants import *
+from sc2.player import Bot, Computer
+
+class WarpGateBot(sc2.BotAI):
+
+    def __init__(self):
+        self.warpgate_started = False
+
+    def select_target(self, state):
+        if self.known_enemy_structures.exists:
+            return random.choice(self.known_enemy_structures)
+
+        return self.enemy_start_locations[0]
+
+    async def on_step(self, iteration):
+        if iteration == 0:
+            await self.chat_send("(glhf)")
+
+        if not self.units(NEXUS).ready.exists:
+            for worker in self.workers:
+                await self.do(worker.attack(self.enemy_start_locations[0]))
+            return
+        else:
+            nexus = self.units(NEXUS).ready.random
+
+        for idle_worker in self.workers.idle:
+            mf = self.state.mineral_field.closest_to(idle_worker)
+            await self.do(idle_worker.gather(mf))
+
+
+        for a in self.units(ASSIMILATOR):
+            if a.assigned_harvesters < a.ideal_harvesters:
+                w = self.workers.closer_than(20, a)
+                if w.exists:
+                    await self.do(w.random.gather(a))
+
+        if self.supply_left < 2 and not self.already_pending(PYLON):
+            if self.can_afford(PYLON):
+                await self.build(PYLON, near=nexus)
+            return
+
+        if self.workers.amount < self.units(NEXUS).amount*15 and nexus.noqueue:
+            if self.can_afford(PROBE):
+                await self.do(nexus.train(PROBE))
+
+        elif not self.units(PYLON).amount < 5 and not self.already_pending(PYLON):
+            if self.can_afford(PYLON):
+                await self.build(PYLON, near=nexus)
+
+        if self.units(NEXUS).amount < 3 and not self.already_pending(NEXUS):
+            if self.can_afford(NEXUS):
+                location = await self.get_next_expansion()
+                await self.build(NEXUS, near=location)
+
+        if self.units(PYLON).ready.exists:
+            pylon = self.units(PYLON).ready.random
+            if self.units(GATEWAY).ready.exists:
+                if not self.units(CYBERNETICSCORE).exists:
+                    if self.can_afford(CYBERNETICSCORE) and not self.already_pending(CYBERNETICSCORE):
+                        await self.build(CYBERNETICSCORE, near=pylon)
+            else:
+                if self.can_afford(GATEWAY) and self.units(GATEWAY).amount < 4:
+                    await self.build(GATEWAY, near=pylon)
+
+        for nexus in self.units(NEXUS).ready:
+            vgs = self.state.vespene_geyser.closer_than(20.0, nexus)
+            for vg in vgs:
+                if not self.can_afford(ASSIMILATOR) or self.units(ASSIMILATOR).ready.exists:
+                    break
+
+                worker = self.select_build_worker(vg.position)
+                if worker is None:
+                    break
+
+                if not self.units(ASSIMILATOR).closer_than(1.0, vg).exists:
+                    await self.do(worker.build(ASSIMILATOR, vg))
+
+        if self.units(CYBERNETICSCORE).ready.exists and self.can_afford(RESEARCH_WARPGATE) and not self.warpgate_started:
+            ccore = self.units(CYBERNETICSCORE).ready.first
+            await self.do(ccore(RESEARCH_WARPGATE))
+            self.warpgate_started = True
+
+        for gateway in self.units(GATEWAY).ready:
+            abilities = await self.get_available_abilities(gateway)
+            if AbilityId.MORPH_WARPGATE in abilities and self.can_afford(AbilityId.MORPH_WARPGATE):
+                await self.do(gateway(MORPH_WARPGATE))
+
+        for warpgate in self.units(WARPGATE).ready:
+            abilities = await self.get_available_abilities(warpgate)
+            # all the units have the same cooldown anyway so let's just look at ZEALOT
+            if AbilityId.TRAINWARP_ZEALOT in abilities:
+                placement = await self.find_placement(AbilityId.TRAINWARP_STALKER, warpgate.position.to2, placement_step=1)
+                if placement is None:
+                    #return ActionResult.CantFindPlacementLocation
+                    print("can't place")
+                    break
+                await self.do(warpgate.warp_in(STALKER, placement))
+
+        if self.units(STALKER).amount > 10 and iteration % 50 == 0:
+            for vr in self.units(STALKER).idle:
+                await self.do(vr.attack(self.select_target(self.state)))
+
+
+def main():
+    sc2.run_game(sc2.maps.get("Abyssal Reef LE"), [
+        Bot(Race.Protoss, WarpGateBot()),
+        Computer(Race.Protoss, Difficulty.Easy)
+    ], realtime=False)
+
+if __name__ == '__main__':
+    main()

sc2/data.py

@@ -12,6 +12,16 @@
 from .ids.unit_typeid import HATCHERY, LAIR, HIVE
 from .ids.unit_typeid import ASSIMILATOR, REFINERY, EXTRACTOR
 
+from .ids.ability_id import TRAIN_ZEALOT, TRAIN_STALKER, TRAIN_HIGHTEMPLAR, TRAIN_DARKTEMPLAR, TRAIN_SENTRY, \
+    TRAIN_ADEPT
+from .ids.ability_id import \
+    TRAINWARP_ZEALOT, \
+    TRAINWARP_STALKER, \
+    TRAINWARP_HIGHTEMPLAR, \
+    TRAINWARP_DARKTEMPLAR, \
+    TRAINWARP_SENTRY, \
+    TRAINWARP_ADEPT
+
 PlayerType = enum.Enum("PlayerType", sc_pb.PlayerType.items())
 Difficulty = enum.Enum("Difficulty", sc_pb.Difficulty.items())
 Status = enum.Enum("Status", sc_pb.Status.items())
@@ -29,7 +39,6 @@
 
 ActionResult = enum.Enum("ActionResult", error_pb.ActionResult.items())
 
-
 race_worker = {
     Race.Protoss: PROBE,
     Race.Terran: SCV,
@@ -42,8 +51,17 @@
     Race.Zerg: {HATCHERY, LAIR, HIVE}
 }
 
+warpgate_abilities = {
+    TRAIN_ZEALOT: TRAINWARP_ZEALOT,
+    TRAIN_STALKER: TRAINWARP_STALKER,
+    TRAIN_HIGHTEMPLAR: TRAINWARP_HIGHTEMPLAR,
+    TRAIN_DARKTEMPLAR: TRAINWARP_DARKTEMPLAR,
+    TRAIN_SENTRY: TRAINWARP_SENTRY,
+    TRAIN_ADEPT: TRAINWARP_ADEPT
+}
+
 race_gas = {
     Race.Protoss: ASSIMILATOR,
     Race.Terran: REFINERY,
     Race.Zerg: EXTRACTOR
-}
+}

sc2/unit.py

@@ -2,7 +2,7 @@
 from sc2.ids.buff_id import BuffId
 
 from .position import Point3
-from .data import Alliance, Attribute, DisplayType
+from .data import Alliance, Attribute, DisplayType, warpgate_abilities
 from .game_data import GameData
 from .ids.unit_typeid import UnitTypeId
 from .ids.ability_id import AbilityId
@@ -190,7 +190,11 @@ def has_buff(self, buff):
 
         return buff.value in self._proto.buff_ids
 
-    def attack(self, *args, **kwargs):
+    def warp_in(self, unit, placement, *args, **kwargs):
+        normal_creation_ability = self._game_data.units[unit.value].creation_ability.id
+        return self(warpgate_abilities[normal_creation_ability], placement, *args, **kwargs)
+
+      def attack(self, *args, **kwargs):
         return self(AbilityId.ATTACK, *args, **kwargs)
 
     def gather(self, *args, **kwargs):