add type stubs for static type checking support

.github/workflows/docker-ci.yml

@@ -111,6 +111,7 @@ jobs:
           docker cp pyproject.toml test_container:/root/python-sc2/
           docker cp uv.lock test_container:/root/python-sc2/
           docker cp sc2 test_container:/root/python-sc2/sc2
+          docker cp s2clientprotocol test_container:/root/python-sc2/s2clientprotocol
           docker cp test test_container:/root/python-sc2/test
           docker cp examples test_container:/root/python-sc2/examples
           docker exec -i test_container bash -c "pip install uv \

dockerfiles/test_docker_image.sh

@@ -46,6 +46,7 @@ docker cp uv.lock test_container:/root/python-sc2/
 docker exec -i test_container bash -c "pip install uv && cd python-sc2 && uv sync --no-cache --no-install-project"
 
 docker cp sc2 test_container:/root/python-sc2/sc2
+docker cp s2clientprotocol test_container:/root/python-sc2/s2clientprotocol
 docker cp test test_container:/root/python-sc2/test
 
 # Run various test bots

examples/arcade_bot.py

@@ -42,7 +42,7 @@ async def on_start(self):
         await self.chat_send("Edit this message for automatic chat commands.")
         self.client.game_step = 2
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         # do marine micro vs zerglings
         for unit in self.units(UnitTypeId.MARINE):
             if self.enemy_units:

examples/competitive/bot.py

@@ -7,7 +7,7 @@ async def on_start(self):
         print("Game started")
         # Do things here before the game starts
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         # Populate this function with whatever your bot should do!
         pass
 

examples/distributed_workers.py

@@ -8,7 +8,7 @@
 
 
 class TerranBot(BotAI):
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         await self.distribute_workers()
         await self.build_supply()
         await self.build_workers()

examples/fastreload.py

@@ -4,11 +4,14 @@
 from sc2 import maps
 from sc2.data import Difficulty, Race
 from sc2.main import _host_game_iter
-from sc2.player import Bot, Computer
+from sc2.player import AbstractPlayer, Bot, Computer
 
 
 def main():
-    player_config = [Bot(Race.Zerg, zerg_rush.ZergRushBot()), Computer(Race.Terran, Difficulty.Medium)]
+    player_config: list[AbstractPlayer] = [
+        Bot(Race.Zerg, zerg_rush.ZergRushBot()),
+        Computer(Race.Terran, Difficulty.Medium),
+    ]
 
     gen = _host_game_iter(maps.get("Abyssal Reef LE"), player_config, realtime=False)
 

examples/host_external_norestart.py

@@ -1,13 +1,13 @@
-import sc2
 from examples.zerg.zerg_rush import ZergRushBot
 from sc2 import maps
 from sc2.data import Race
 from sc2.main import _host_game_iter
 from sc2.player import Bot
+from sc2.portconfig import Portconfig
 
 
 def main():
-    portconfig = sc2.portconfig.Portconfig()
+    portconfig: Portconfig = Portconfig()
     print(portconfig.as_json)
 
     player_config = [Bot(Race.Zerg, ZergRushBot()), Bot(Race.Zerg, None)]

examples/protoss/cannon_rush.py

@@ -9,7 +9,7 @@
 
 
 class CannonRushBot(BotAI):
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         if iteration == 0:
             await self.chat_send("(probe)(pylon)(cannon)(cannon)(gg)")
 

examples/protoss/find_adept_shades.py

@@ -13,7 +13,7 @@
 class FindAdeptShadesBot(BotAI):
     def __init__(self):
         self.shaded = False
-        self.shades_mapping = {}
+        self.shades_mapping: dict[int, int] = {}
 
     async def on_start(self):
         self.client.game_step = 2

examples/protoss/threebase_voidray.py

@@ -9,7 +9,7 @@
 
 
 class ThreebaseVoidrayBot(BotAI):
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         target_base_count = 3
         target_stargate_count = 3
 

examples/protoss/warpgate_push.py

@@ -9,18 +9,19 @@
 from sc2.ids.upgrade_id import UpgradeId
 from sc2.main import run_game
 from sc2.player import Bot, Computer
+from sc2.unit import Unit
 
 
 class WarpGateBot(BotAI):
     def __init__(self):
         # Initialize inherited class
         self.proxy_built = False
 
-    async def warp_new_units(self, proxy):
+    async def warp_new_units(self, proxy: Unit):
         for warpgate in self.structures(UnitTypeId.WARPGATE).ready:
-            abilities = await self.get_available_abilities(warpgate)
+            abilities = await self.get_available_abilities([warpgate])
             # all the units have the same cooldown anyway so let's just look at ZEALOT
-            if AbilityId.WARPGATETRAIN_STALKER in abilities:
+            if AbilityId.WARPGATETRAIN_STALKER in abilities[0]:
                 pos = proxy.position.to2.random_on_distance(4)
                 placement = await self.find_placement(AbilityId.WARPGATETRAIN_STALKER, pos, placement_step=1)
                 if placement is None:
@@ -29,7 +30,7 @@ async def warp_new_units(self, proxy):
                     return
                 warpgate.warp_in(UnitTypeId.STALKER, placement)
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         await self.distribute_workers()
 
         if not self.townhalls.ready:

examples/simulate_fight_scenario.py

@@ -15,15 +15,15 @@
 class FightBot(BotAI):
     def __init__(self):
         super().__init__()
-        self.enemy_location: Point2 = None
+        self.enemy_location: Point2 | None = None
         self.fight_started = False
 
     async def on_start(self):
         # Retrieve control by enabling enemy control and showing whole map
         await self.client.debug_show_map()
         await self.client.debug_control_enemy()
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         # Wait till control retrieved, destroy all starting units, recreate the world
         if iteration > 0 and self.enemy_units and not self.enemy_location:
             await self.reset_arena()

examples/terran/cyclone_push.py

@@ -28,7 +28,7 @@ def select_target(self) -> Point2:
         # Pick a random mineral field on the map
         return self.mineral_field.random.position
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         CCs: Units = self.townhalls(UnitTypeId.COMMANDCENTER)
         # If no command center exists, attack-move with all workers and cyclones
         if not CCs:
@@ -87,7 +87,7 @@ async def on_step(self, iteration):
                         if self.gas_buildings.filter(lambda unit: unit.distance_to(vg) < 1):
                             continue
                         # Select a worker closest to the vespene geysir
-                        worker: Unit = self.select_build_worker(vg)
+                        worker: Unit | None = self.select_build_worker(vg)
                         # Worker can be none in cases where all workers are dead
                         # or 'select_build_worker' function only selects from workers which carry no minerals
                         if worker is None:
@@ -112,9 +112,9 @@ async def on_step(self, iteration):
         # Saturate gas
         for refinery in self.gas_buildings:
             if refinery.assigned_harvesters < refinery.ideal_harvesters:
-                worker: Units = self.workers.closer_than(10, refinery)
-                if worker:
-                    worker.random.gather(refinery)
+                workers: Units = self.workers.closer_than(10, refinery)
+                if workers:
+                    workers.random.gather(refinery)
 
         for scv in self.workers.idle:
             scv.gather(self.mineral_field.closest_to(cc))

examples/terran/mass_reaper.py

@@ -25,7 +25,7 @@ def __init__(self):
         # Select distance calculation method 0, which is the pure python distance calculation without caching or indexing, using math.hypot(), for more info see bot_ai_internal.py _distances_override_functions() function
         self.distance_calculation_method = 3
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         # Benchmark and print duration time of the on_step method based on "self.distance_calculation_method" value
         # logger.info(self.time_formatted, self.supply_used, self.step_time[1])
         """
@@ -45,7 +45,9 @@ async def on_step(self, iteration):
             # If workers were found
             if workers:
                 worker: Unit = workers.furthest_to(workers.center)
-                location: Point2 = await self.find_placement(UnitTypeId.SUPPLYDEPOT, worker.position, placement_step=3)
+                location: Point2 | None = await self.find_placement(
+                    UnitTypeId.SUPPLYDEPOT, worker.position, placement_step=3
+                )
                 # If a placement location was found
                 if location:
                     # Order worker to build exactly on that location
@@ -72,13 +74,13 @@ async def on_step(self, iteration):
             and self.can_afford(UnitTypeId.COMMANDCENTER)
         ):
             # get_next_expansion returns the position of the next possible expansion location where you can place a command center
-            location: Point2 = await self.get_next_expansion()
+            location: Point2 | None = await self.get_next_expansion()
             if location:
                 # Now we "select" (or choose) the nearest worker to that found location
-                worker: Unit = self.select_build_worker(location)
-                if worker and self.can_afford(UnitTypeId.COMMANDCENTER):
+                worker2: Unit | None = self.select_build_worker(location)
+                if worker2 and self.can_afford(UnitTypeId.COMMANDCENTER):
                     # The worker will be commanded to build the command center
-                    worker.build(UnitTypeId.COMMANDCENTER, location)
+                    worker2.build(UnitTypeId.COMMANDCENTER, location)
 
         # Build up to 4 barracks if we can afford them
         # Check if we have a supply depot (tech requirement) before trying to make barracks
@@ -97,7 +99,7 @@ async def on_step(self, iteration):
             ):  # need to check if townhalls.amount > 0 because placement is based on townhall location
                 worker: Unit = workers.furthest_to(workers.center)
                 # I chose placement_step 4 here so there will be gaps between barracks hopefully
-                location: Point2 = await self.find_placement(
+                location: Point2 | None = await self.find_placement(
                     UnitTypeId.BARRACKS, self.townhalls.random.position, placement_step=4
                 )
                 if location:
@@ -168,7 +170,7 @@ async def on_step(self, iteration):
                 retreat_points: set[Point2] = {x for x in retreat_points if self.in_pathing_grid(x)}
                 if retreat_points:
                     closest_enemy: Unit = enemy_threats_close.closest_to(r)
-                    retreat_point: Unit = closest_enemy.position.furthest(retreat_points)
+                    retreat_point: Point2 = closest_enemy.position.furthest(retreat_points)
                     r.move(retreat_point)
                     continue  # Continue for loop, dont execute any of the following
 
@@ -259,13 +261,13 @@ async def on_step(self, iteration):
     # Stolen and modified from position.py
 
     @staticmethod
-    def neighbors4(position, distance=1) -> set[Point2]:
+    def neighbors4(position: Point2, distance: float = 1) -> set[Point2]:
         p = position
         d = distance
         return {Point2((p.x - d, p.y)), Point2((p.x + d, p.y)), Point2((p.x, p.y - d)), Point2((p.x, p.y + d))}
 
     # Stolen and modified from position.py
-    def neighbors8(self, position, distance=1) -> set[Point2]:
+    def neighbors8(self, position: Point2, distance: float = 1) -> set[Point2]:
         p = position
         d = distance
         return self.neighbors4(position, distance) | {

examples/terran/onebase_battlecruiser.py

@@ -29,7 +29,7 @@ def select_target(self) -> tuple[Point2, bool]:
 
         return self.mineral_field.random.position, False
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         ccs: Units = self.townhalls
         # If we no longer have townhalls, attack with all workers
         if not ccs:
@@ -85,7 +85,7 @@ async def on_step(self, iteration):
                         if self.gas_buildings.filter(lambda unit: unit.distance_to(vg) < 1):
                             break
 
-                        worker: Unit = self.select_build_worker(vg.position)
+                        worker: Unit | None = self.select_build_worker(vg.position)
                         if worker is None:
                             break
 
@@ -172,9 +172,9 @@ def starport_land_positions(sp_position: Point2) -> list[Point2]:
         # Saturate refineries
         for refinery in self.gas_buildings:
             if refinery.assigned_harvesters < refinery.ideal_harvesters:
-                worker: Units = self.workers.closer_than(10, refinery)
-                if worker:
-                    worker.random.gather(refinery)
+                workers: Units = self.workers.closer_than(10, refinery)
+                if workers:
+                    workers.random.gather(refinery)
 
         # Send workers back to mine if they are idle
         for scv in self.workers.idle:

examples/terran/proxy_rax.py

@@ -13,7 +13,7 @@ class ProxyRaxBot(BotAI):
     async def on_start(self):
         self.client.game_step = 2
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         # If we don't have a townhall anymore, send all units to attack
         ccs: Units = self.townhalls(UnitTypeId.COMMANDCENTER)
         if not ccs:

examples/terran/ramp_wall.py

@@ -19,7 +19,7 @@ class RampWallBot(BotAI):
     def __init__(self):
         self.unit_command_uses_self_do = False
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         ccs: Units = self.townhalls(UnitTypeId.COMMANDCENTER)
         if not ccs:
             return
@@ -70,11 +70,11 @@ async def on_step(self, iteration):
         # Draw if two selected units are facing each other - green if this guy is facing the other, red if he is not
         self.draw_facing_units()
 
-        depot_placement_positions: frozenset[Point2] = self.main_base_ramp.corner_depots
+        depot_placement_positions: set[Point2] = self.main_base_ramp.corner_depots
         # Uncomment the following if you want to build 3 supply depots in the wall instead of a barracks in the middle + 2 depots in the corner
         # depot_placement_positions = self.main_base_ramp.corner_depots | {self.main_base_ramp.depot_in_middle}
 
-        barracks_placement_position: Point2 = self.main_base_ramp.barracks_correct_placement
+        barracks_placement_position: Point2 | None = self.main_base_ramp.barracks_correct_placement
         # If you prefer to have the barracks in the middle without room for addons, use the following instead
         # barracks_placement_position = self.main_base_ramp.barracks_in_middle
 

examples/too_slow_bot.py

@@ -9,7 +9,7 @@
 
 
 class SlowBot(ProxyRaxBot):
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         await asyncio.sleep(random.random())
         await super().on_step(iteration)
 

examples/worker_rush.py

@@ -6,7 +6,7 @@
 
 
 class WorkerRushBot(BotAI):
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         if iteration == 0:
             for worker in self.workers:
                 worker.attack(self.enemy_start_locations[0])

examples/zerg/banes_banes_banes.py

@@ -23,7 +23,7 @@ def select_target(self) -> Point2:
             return random.choice(self.enemy_structures).position
         return self.enemy_start_locations[0]
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         larvae: Units = self.larva
         lings: Units = self.units(UnitTypeId.ZERGLING)
         # Send all idle banes to enemy

examples/zerg/expand_everywhere.py

@@ -16,7 +16,7 @@ async def on_start(self):
         self.client.game_step = 50
         await self.client.debug_show_map()
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         # Build overlords if about to be supply blocked
         if (
             self.supply_left < 2

examples/zerg/hydralisk_push.py

@@ -19,7 +19,7 @@ def select_target(self) -> Point2:
             return random.choice(self.enemy_structures).position
         return self.enemy_start_locations[0]
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         larvae: Units = self.larva
         forces: Units = self.units.of_type({UnitTypeId.ZERGLING, UnitTypeId.HYDRALISK})
 

examples/zerg/onebase_broodlord.py

@@ -19,7 +19,7 @@ def select_target(self) -> Point2:
             return random.choice(self.enemy_structures).position
         return self.enemy_start_locations[0]
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         larvae: Units = self.larva
         forces: Units = self.units.of_type({UnitTypeId.ZERGLING, UnitTypeId.CORRUPTOR, UnitTypeId.BROODLORD})
 

examples/zerg/worker_split.py

@@ -30,7 +30,7 @@ async def on_before_start(self):
     async def on_start(self):
         """This function is run after the expansion locations and ramps are calculated."""
 
-    async def on_step(self, iteration):
+    async def on_step(self, iteration: int):
         if iteration % 10 == 0:
             await asyncio.sleep(3)
         # In realtime=False, this should print "8*x" and "x" if