Fix types for bot_ai.py

Author: BurnySc2

pyproject.toml

@@ -101,8 +101,6 @@ project_includes = [
 project-excludes = [
     # Disable for those files and folders
     "sc2/data.py",
-    # TODO Temp disable for those files and folders
-    "sc2/bot_ai.py",
 ]
 
 [tool.pyrefly.errors]

sc2/bot_ai.py

@@ -250,24 +250,24 @@ async def get_next_expansion(self) -> Point2 | None:
         """Find next expansion location."""
 
         closest = None
-        distance = math.inf
-        for el in self.expansion_locations_list:
+        best_distance = math.inf
+        start_position = self.game_info.player_start_location
+        for position in self.expansion_locations_list:
 
             def is_near_to_expansion(t):
-                return t.distance_to(el) < self.EXPANSION_GAP_THRESHOLD
+                return t.distance_to(position) < self.EXPANSION_GAP_THRESHOLD
 
             if any(map(is_near_to_expansion, self.townhalls)):
                 # already taken
                 continue
 
-            startp = self.game_info.player_start_location
-            d = await self.client.query_pathing(startp, el)
-            if d is None:
+            distance = await self.client.query_pathing(start_position, position)
+            if distance is None:
                 continue
 
-            if d < distance:
-                distance = d
-                closest = el
+            if distance < best_distance:
+                best_distance = distance
+                closest = position
 
         return closest
 
@@ -492,7 +492,11 @@ def calculate_cost(self, item_id: UnitTypeId | UpgradeId | AbilityId) -> Cost:
                     return self.calculate_unit_value(UnitTypeId.ARCHON)
             unit_data = self.game_data.units[item_id.value]
             # Cost of morphs is automatically correctly calculated by 'calculate_ability_cost'
-            return self.game_data.calculate_ability_cost(unit_data.creation_ability.exact_id)
+            creation_ability = unit_data.creation_ability
+            if creation_ability is None:
+                logger.error(f"Unknown item_id for calculate_cost: {item_id}")
+                return Cost(0, 0)
+            return self.game_data.calculate_ability_cost(creation_ability.exact_id)
 
         if isinstance(item_id, UpgradeId):
             cost = self.game_data.upgrades[item_id.value].cost
@@ -566,11 +570,18 @@ async def can_cast(
             ability_target: int = self.game_data.abilities[ability_id.value]._proto.target
             # Check if target is in range (or is a self cast like stimpack)
             if (
+                # Can't replace 1 with "Target.None.value" because ".None" doesn't seem to be a valid enum name
                 ability_target == 1
                 or ability_target == Target.PointOrNone.value
-                and isinstance(target, Point2)
-                and unit.distance_to(target) <= unit.radius + target.radius + cast_range
-            ):  # cant replace 1 with "Target.None.value" because ".None" doesnt seem to be a valid enum name
+                and (
+                    # Target is unit
+                    isinstance(target, Unit)
+                    and unit.distance_to(target) <= unit.radius + target.radius + cast_range
+                    # Target is position
+                    or isinstance(target, Point2)
+                    and unit.distance_to(target) <= unit.radius + cast_range
+                )
+            ):
                 return True
             # Check if able to use ability on a unit
             if (
@@ -620,8 +631,12 @@ def select_build_worker(self, pos: Unit | Point2, force: bool = False) -> Unit |
     async def can_place_single(self, building: AbilityId | UnitTypeId, position: Point2) -> bool:
         """Checks the placement for only one position."""
         if isinstance(building, UnitTypeId):
-            creation_ability = self.game_data.units[building.value].creation_ability.id
-            return (await self.client._query_building_placement_fast(creation_ability, [position]))[0]
+            creation_ability = self.game_data.units[building.value].creation_ability
+            if creation_ability is None:
+                logger.error(f"Unknown building for can_place_single: {building}")
+                return False
+            creation_ability_id = creation_ability.id
+            return (await self.client._query_building_placement_fast(creation_ability_id, [position]))[0]
         return (await self.client._query_building_placement_fast(building, [position]))[0]
 
     async def can_place(self, building: AbilityData | AbilityId | UnitTypeId, positions: list[Point2]) -> list[bool]:
@@ -637,24 +652,26 @@ async def can_place(self, building: AbilityData | AbilityId | UnitTypeId, positi
 
         :param building:
         :param position:"""
-        building_type = type(building)
-        assert type(building) in {AbilityData, AbilityId, UnitTypeId}, f"{building}, {building_type}"
-        if building_type == UnitTypeId:
-            building = self.game_data.units[building.value].creation_ability.id
-        elif building_type == AbilityData:
+        if isinstance(building, UnitTypeId):
+            creation_ability = self.game_data.units[building.value].creation_ability
+            if creation_ability is None:
+                return [False for _ in positions]
+            building = creation_ability.id
+        elif isinstance(building, AbilityData):
             warnings.warn(
                 "Using AbilityData is deprecated and may be removed soon. Please use AbilityId or UnitTypeId instead.",
                 DeprecationWarning,
                 stacklevel=2,
             )
-            building = building_type.id
+            building = building.id
 
         if isinstance(positions, (Point2, tuple)):
             warnings.warn(
                 "The support for querying single entries will be removed soon. Please use either 'await self.can_place_single(building, position)' or 'await (self.can_place(building, [position]))[0]",
                 DeprecationWarning,
                 stacklevel=2,
             )
+            # pyrefly: ignore
             return await self.can_place_single(building, positions)
         assert isinstance(positions, list), f"Expected an iterable (list, tuple), but was: {positions}"
         assert isinstance(positions[0], Point2), (
@@ -690,7 +707,10 @@ async def find_placement(
         assert isinstance(near, Point2), f"{near} is no Point2 object"
 
         if isinstance(building, UnitTypeId):
-            building = self.game_data.units[building.value].creation_ability.id
+            creation_ability = self.game_data.units[building.value].creation_ability
+            if creation_ability is None:
+                return None
+            building = creation_ability.id
 
         if await self.can_place_single(building, near) and (
             not addon_place or await self.can_place_single(AbilityId.TERRANBUILD_SUPPLYDEPOT, near.offset((2.5, -0.5)))
@@ -749,10 +769,13 @@ def already_pending_upgrade(self, upgrade_type: UpgradeId) -> float:
         assert isinstance(upgrade_type, UpgradeId), f"{upgrade_type} is no UpgradeId"
         if upgrade_type in self.state.upgrades:
             return 1
-        creationAbilityID = self.game_data.upgrades[upgrade_type.value].research_ability.exact_id
+        research_ability = self.game_data.upgrades[upgrade_type.value].research_ability
+        if research_ability is None:
+            return 0
+        creation_ability_id = research_ability.exact_id
         for structure in self.structures.filter(lambda unit: unit.is_ready):
             for order in structure.orders:
-                if order.ability.exact_id == creationAbilityID:
+                if order.ability.exact_id == creation_ability_id:
                     return order.progress
         return 0
 
@@ -798,7 +821,7 @@ def structure_type_build_progress(self, structure_type: UnitTypeId | int) -> flo
         # SUPPLYDEPOTDROP is not in self.game_data.units, so bot_ai should not check the build progress via creation ability (worker abilities)
         if structure_type_value not in self.game_data.units:
             return max((s.build_progress for s in self.structures if s._proto.unit_type in equiv_values), default=0)
-        creation_ability_data: AbilityData = self.game_data.units[structure_type_value].creation_ability
+        creation_ability_data = self.game_data.units[structure_type_value].creation_ability
         if creation_ability_data is None:
             return 0
         creation_ability: AbilityId = creation_ability_data.exact_id
@@ -863,24 +886,33 @@ def already_pending(self, unit_type: UpgradeId | UnitTypeId) -> float:
         """
         if isinstance(unit_type, UpgradeId):
             return self.already_pending_upgrade(unit_type)
-        try:
-            ability = self.game_data.units[unit_type.value].creation_ability.exact_id
-        except AttributeError:
-            if unit_type in CREATION_ABILITY_FIX:
-                # Hotfix for checking pending archons
-                if unit_type == UnitTypeId.ARCHON:
-                    return self._abilities_count_and_build_progress[0][AbilityId.ARCHON_WARP_TARGET] / 2
-                # Hotfix for rich geysirs
-                return self._abilities_count_and_build_progress[0][CREATION_ABILITY_FIX[unit_type]]
-            logger.error(f"Uncaught UnitTypeId: {unit_type}")
+        creation_ability = self.game_data.units[unit_type.value].creation_ability
+        if creation_ability is None:
             return 0
-        return self._abilities_count_and_build_progress[0][ability]
+
+        if unit_type in CREATION_ABILITY_FIX:
+            # Hotfix for checking pending archons and other abilities
+            if unit_type == UnitTypeId.ARCHON:
+                return self._abilities_count_and_build_progress[0][AbilityId.ARCHON_WARP_TARGET] / 2
+            # Hotfix for rich geysirs
+            return self._abilities_count_and_build_progress[0][CREATION_ABILITY_FIX[unit_type]]
+
+        creation_ability = self.game_data.units[unit_type.value].creation_ability
+        if creation_ability is None:
+            logger.error(f"Unknown unit_type for already_pending: {unit_type}")
+            return 0
+        ability_id = creation_ability.exact_id
+        return self._abilities_count_and_build_progress[0][ability_id]
 
     def worker_en_route_to_build(self, unit_type: UnitTypeId) -> float:
         """This function counts how many workers are on the way to start the construction a building.
 
         :param unit_type:"""
-        ability = self.game_data.units[unit_type.value].creation_ability.exact_id
+        creation_ability = self.game_data.units[unit_type.value].creation_ability
+        if creation_ability is None:
+            logger.error(f"Unknown unit_type for worker_en_route_to_build: {unit_type}")
+            return 0
+        ability = creation_ability.exact_id
         return self._worker_orders[ability]
 
     @property_cache_once_per_frame
@@ -894,7 +926,7 @@ def structures_without_construction_SCVs(self) -> Units:
                 continue
             for order in worker.orders:
                 # When a construction is resumed, the worker.orders[0].target is the tag of the structure, else it is a Point2
-                worker_targets.add(order.target)
+                worker_targets.add(order.target)  # pyrefly: ignore
         return self.structures.filter(
             lambda structure: structure.build_progress < 1
             # Redundant check?
@@ -928,15 +960,15 @@ async def build(
         assert isinstance(near, (Unit, Point2))
         if not self.can_afford(building):
             return False
-        p = None
+        position = None
         gas_buildings = {UnitTypeId.EXTRACTOR, UnitTypeId.ASSIMILATOR, UnitTypeId.REFINERY}
         if isinstance(near, Unit) and building not in gas_buildings:
             near = near.position
         if isinstance(near, Point2):
             near = near.to2
         if isinstance(near, Point2):
-            p = await self.find_placement(building, near, max_distance, random_alternative, placement_step)
-            if p is None:
+            position = await self.find_placement(building, near, max_distance, random_alternative, placement_step)
+            if position is None:
                 return False
         builder = build_worker or self.select_build_worker(near)
         if builder is None:
@@ -945,7 +977,8 @@ async def build(
             assert isinstance(near, Unit)
             builder.build_gas(near)
             return True
-        self.do(builder.build(building, p), subtract_cost=True, ignore_warning=True)
+        # pyrefly: ignore
+        self.do(builder.build(building, position), subtract_cost=True, ignore_warning=True)
         return True
 
     def train(
@@ -1054,6 +1087,7 @@ def train(
                 else:
                     # Normal train a unit from larva or inside a structure
                     successfully_trained = self.do(
+                        # pyrefly: ignore
                         structure.train(unit_type),
                         subtract_cost=True,
                         subtract_supply=True,
@@ -1077,6 +1111,7 @@ def train(
                         trained_amount += 1
                         # With one command queue=False and one queue=True, you can queue 2 marines in a reactored barracks in one frame
                         successfully_trained = self.do(
+                            # pyrefly: ignore
                             structure.train(unit_type, queue=True),
                             subtract_cost=True,
                             subtract_supply=True,
@@ -1124,6 +1159,7 @@ def research(self, upgrade_type: UpgradeId) -> bool:
 
         research_structure_type: UnitTypeId = UPGRADE_RESEARCHED_FROM[upgrade_type]
 
+        # pyrefly: ignore
         required_tech_building: UnitTypeId | None = RESEARCH_INFO[research_structure_type][upgrade_type].get(
             "required_building", None
         )
@@ -1166,6 +1202,7 @@ def research(self, upgrade_type: UpgradeId) -> bool:
             ):
                 # Can_afford check was already done earlier in this function
                 successful_action: bool = self.do(
+                    # pyrefly: ignore
                     structure.research(upgrade_type),
                     subtract_cost=True,
                     ignore_warning=True,

sc2/client.py

@@ -226,7 +226,7 @@ async def actions(self, actions: list[UnitCommand], return_successes: bool = Fal
             ActionResult(result) for result in response.action.result if ActionResult(result) != ActionResult.Success
         ]
 
-    async def query_pathing(self, start: Unit | Point2 | Point3, end: Point2 | Point3) -> int | float | None:
+    async def query_pathing(self, start: Unit | Point2 | Point3, end: Point2 | Point3) -> float | None:
         """Caution: returns "None" when path not found
         Try to combine queries with the function below because the pathing query is generally slow.