Merge pull request #176 from BurnySc2/fix_creation_ability_error

Author: BurnySc2

sc2/bot_ai.py

@@ -13,6 +13,7 @@
 from sc2.bot_ai_internal import BotAIInternal
 from sc2.cache import property_cache_once_per_frame
 from sc2.constants import (
+    CREATION_ABILITY_FIX,
     EQUIVALENTS_FOR_TECH_PROGRESS,
     PROTOSS_TECH_REQUIREMENT,
     TERRAN_STRUCTURES_REQUIRE_SCV,
@@ -791,7 +792,7 @@ def structure_type_build_progress(self, structure_type: Union[UnitTypeId, int])
         creation_ability: AbilityId = creation_ability_data.exact_id
         max_value = max(
             [s.build_progress for s in self.structures if s._proto.unit_type in equiv_values] +
-            [self._abilities_all_units[1].get(creation_ability, 0)],
+            [self._abilities_count_and_build_progress[1].get(creation_ability, 0)],
             default=0,
         )
         return max_value
@@ -853,12 +854,15 @@ def already_pending(self, unit_type: Union[UpgradeId, UnitTypeId]) -> float:
         try:
             ability = self.game_data.units[unit_type.value].creation_ability.exact_id
         except AttributeError:
-            # Hotfix for checking pending archons
-            if unit_type == UnitTypeId.ARCHON:
-                return self._abilities_all_units[0][AbilityId.ARCHON_WARP_TARGET] / 2
+            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}")
             return 0
-        return self._abilities_all_units[0][ability]
+        return self._abilities_count_and_build_progress[0][ability]
 
     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.

sc2/bot_ai_internal.py

@@ -19,6 +19,7 @@
 from sc2.cache import property_cache_once_per_frame
 from sc2.constants import (
     ALL_GAS,
+    CREATION_ABILITY_FIX,
     IS_PLACEHOLDER,
     TERRAN_STRUCTURES_REQUIRE_SCV,
     FakeEffectID,
@@ -266,7 +267,7 @@ def _correct_zerg_supply(self):
 
     @final
     @property_cache_once_per_frame
-    def _abilities_all_units(self) -> Tuple[CounterType[AbilityId], Dict[AbilityId, float]]:
+    def _abilities_count_and_build_progress(self) -> Tuple[CounterType[AbilityId], Dict[AbilityId, float]]:
         """Cache for the already_pending function, includes protoss units warping in,
         all units in production and all structures, and all morphs"""
         abilities_amount: CounterType[AbilityId] = Counter()
@@ -279,10 +280,15 @@ def _abilities_all_units(self) -> Tuple[CounterType[AbilityId], Dict[AbilityId,
                 if self.race != Race.Terran or not unit.is_structure:
                     # If an SCV is constructing a building, already_pending would count this structure twice
                     # (once from the SCV order, and once from "not structure.is_ready")
-                    if unit.type_id == UnitTypeId.ARCHON:
-                        # Hotfix for archons in morph state
-                        creation_ability = AbilityId.ARCHON_WARP_TARGET
-                        abilities_amount[creation_ability] += 2
+                    if unit.type_id in CREATION_ABILITY_FIX:
+                        if unit.type_id == UnitTypeId.ARCHON:
+                            # Hotfix for archons in morph state
+                            creation_ability = AbilityId.ARCHON_WARP_TARGET
+                            abilities_amount[creation_ability] += 2
+                        else:
+                            # Hotfix for rich geysirs
+                            creation_ability = CREATION_ABILITY_FIX[unit.type_id]
+                            abilities_amount[creation_ability] += 1
                     else:
                         creation_ability: AbilityId = self.game_data.units[unit.type_id.value].creation_ability.exact_id
                         abilities_amount[creation_ability] += 1
@@ -743,7 +749,7 @@ async def _issue_vision_events(self):
         enemy_units_left_vision: Set[int] = set(self._enemy_units_previous_map) - self.enemy_units.tags
         for enemy_unit_tag in enemy_units_left_vision:
             await self.on_enemy_unit_left_vision(enemy_unit_tag)
-        enemy_structures_left_vision: Set[int] = (set(self._enemy_structures_previous_map) - self.enemy_structures.tags)
+        enemy_structures_left_vision: Set[int] = set(self._enemy_structures_previous_map) - self.enemy_structures.tags
         for enemy_structure_tag in enemy_structures_left_vision:
             await self.on_enemy_unit_left_vision(enemy_structure_tag)
 

sc2/constants.py

@@ -681,3 +681,16 @@ def return_NOTAUNIT() -> UnitTypeId:
     4: "Point2 or Unit",
     5: "Point2 or no target",
 }
+CREATION_ABILITY_FIX: Dict[UnitTypeId, AbilityId] = {
+    UnitTypeId.ARCHON: AbilityId.ARCHON_WARP_TARGET,
+    UnitTypeId.ASSIMILATORRICH: AbilityId.PROTOSSBUILD_ASSIMILATOR,
+    UnitTypeId.BANELINGCOCOON: AbilityId.MORPHZERGLINGTOBANELING_BANELING,
+    UnitTypeId.CHANGELING: AbilityId.SPAWNCHANGELING_SPAWNCHANGELING,
+    UnitTypeId.EXTRACTORRICH: AbilityId.ZERGBUILD_EXTRACTOR,
+    UnitTypeId.INTERCEPTOR: AbilityId.BUILD_INTERCEPTORS,
+    UnitTypeId.LURKERMPEGG: AbilityId.MORPH_LURKER,
+    UnitTypeId.MULE: AbilityId.CALLDOWNMULE_CALLDOWNMULE,
+    UnitTypeId.RAVAGERCOCOON: AbilityId.MORPHTORAVAGER_RAVAGER,
+    UnitTypeId.REFINERYRICH: AbilityId.TERRANBUILD_REFINERY,
+    UnitTypeId.TECHLAB: AbilityId.BUILD_TECHLAB,
+}

test/autotest_bot.py

@@ -74,7 +74,7 @@ async def on_step(self, iteration):
     async def clean_up_center(self):
         map_center = self.game_info.map_center
         # Remove everything close to map center
-        my_units = self.units | self.structures
+        my_units = self.all_own_units
         if my_units:
             my_units = my_units.closer_than(20, map_center)
         if my_units:
@@ -143,8 +143,8 @@ async def test_game_info_static_variables(self):
         assert len(self.game_info.player_races) == 2, self.game_info.player_races
         self.tests_done_by_name.add("test_game_info_static_variables")
 
-    # Test BotAI action: train SCV
     async def test_botai_actions1(self):
+        # Test BotAI action: train SCV
         while self.already_pending(UnitTypeId.SCV) < 1:
             if self.can_afford(UnitTypeId.SCV):
                 self.townhalls.random.train(UnitTypeId.SCV)
@@ -184,8 +184,8 @@ def temp_filter(unit: Unit):
         await self._advance_steps(2)
         logger.warning("Action test 02 successful.")
 
-    # Test BotAI action: move some scvs to the center, some to minerals
     async def test_botai_actions3(self):
+        # Test BotAI action: move some scvs to the center, some to minerals
         center = self.game_info.map_center
 
         while self.units.filter(lambda x: x.is_moving).amount < 6 and self.units.gathering.amount >= 6:
@@ -202,8 +202,8 @@ async def test_botai_actions3(self):
         await self._advance_steps(2)
         logger.warning("Action test 03 successful.")
 
-    # Test BotAI action: move all SCVs to mine minerals near townhall
     async def test_botai_actions4(self):
+        # Test BotAI action: move all SCVs to mine minerals near townhall
         while self.units.gathering.amount < 12:
             mf = self.mineral_field.closest_to(self.townhalls.random)
             for scv in self.workers:
@@ -213,8 +213,8 @@ async def test_botai_actions4(self):
         await self._advance_steps(2)
         logger.warning("Action test 04 successful.")
 
-    # Test BotAI action: self.expand_now() which tests for get_next_expansion, select_build_worker, can_place, find_placement, build and can_afford
     async def test_botai_actions5(self):
+        # Test BotAI action: self.expand_now() which tests for get_next_expansion, select_build_worker, can_place, find_placement, build and can_afford
         # Wait till worker has started construction of CC
         while 1:
             if self.can_afford(UnitTypeId.COMMANDCENTER):
@@ -237,8 +237,8 @@ async def test_botai_actions5(self):
         await self._advance_steps(2)
         logger.warning("Action test 05 successful.")
 
-    # Test if reaper grenade shows up in effects
     async def test_botai_actions6(self):
+        # Test if reaper grenade shows up in effects
         center = self.game_info.map_center
 
         while 1:
@@ -263,8 +263,8 @@ async def test_botai_actions6(self):
         await self._advance_steps(100)
         logger.warning("Action test 06 successful.")
 
-    # Test ravager effects
     async def test_botai_actions7(self):
+        # Test ravager effects
         center = self.game_info.map_center
         while 1:
             if self.units(UnitTypeId.RAVAGER).amount < 10:
@@ -287,8 +287,8 @@ async def test_botai_actions7(self):
         await self._advance_steps(100)
         logger.warning("Action test 07 successful.")
 
-    # Test if train function works on hatchery, lair, hive
     async def test_botai_actions8(self):
+        # Test if train function works on hatchery, lair, hive
         center = self.game_info.map_center
         if not self.structures(UnitTypeId.HIVE):
             await self.client.debug_create_unit([[UnitTypeId.HIVE, 1, center, 1]])
@@ -319,8 +319,8 @@ async def test_botai_actions8(self):
         await self._advance_steps(2)
         logger.warning("Action test 08 successful.")
 
-    # Morph an archon from 2 high templars
     async def test_botai_actions9(self):
+        # Morph an archon from 2 high templars
         center = self.game_info.map_center
         await self.client.debug_create_unit(
             [
@@ -360,8 +360,8 @@ async def test_botai_actions9(self):
         await self._advance_steps(2)
         logger.warning("Action test 09 successful.")
 
-    # Morph 400 banelings from 400 lings in the same frame
     async def test_botai_actions10(self):
+        # Morph 400 banelings from 400 lings in the same frame
         center = self.game_info.map_center
 
         target_amount = 400
@@ -381,8 +381,9 @@ async def test_botai_actions10(self):
             # Spawn units
             if not bane_nests:
                 await self.client.debug_create_unit([[UnitTypeId.BANELINGNEST, 1, center, 1]])
-            if banes.amount + bane_cocoons.amount + lings.amount < target_amount:
-                await self.client.debug_create_unit([[UnitTypeId.ZERGLING, target_amount - lings.amount, center, 1]])
+            current_amount = banes.amount + bane_cocoons.amount + lings.amount
+            if current_amount < target_amount:
+                await self.client.debug_create_unit([[UnitTypeId.ZERGLING, target_amount - current_amount, center, 1]])
 
             if lings.amount >= target_amount and self.minerals >= 10_000 and self.vespene >= 10_000:
                 for ling in lings:
@@ -402,8 +403,8 @@ async def test_botai_actions10(self):
         await self._advance_steps(2)
         logger.warning("Action test 10 successful.")
 
-    # Trigger anti armor missile of raven against enemy unit and check if buff was received
     async def test_botai_actions11(self):
+        # Trigger anti armor missile of raven against enemy unit and check if buff was received
         await self.clean_up_center()
         await self.clean_up_center()
 
@@ -435,10 +436,8 @@ async def test_botai_actions11(self):
         logger.warning("Action test 11 successful.")
         await self.clean_up_center()
 
-    # Test if structures_without_construction_SCVs works after killing the scv
     async def test_botai_actions12(self):
-        map_center: Point2 = self.game_info.map_center
-
+        # Test if structures_without_construction_SCVs works after killing the scv
         # Wait till can afford depot
         while not self.can_afford(UnitTypeId.SUPPLYDEPOT):
             await self.client.debug_all_resources()

test/generate_pickle_files_bot.py

@@ -109,6 +109,7 @@ def main():
         "AcidPlantLE",
         "AcolyteLE",
         "AcropolisLE",
+        "AncientCisternAIE",
         "Artana",
         "AscensiontoAiurLE",
         "AutomatonLE",
@@ -132,6 +133,7 @@ def main():
         "DefendersLandingLE",
         "DigitalFrontier",
         "DiscoBloodbathLE",
+        "DragonScalesAIE",
         "DreamcatcherLE",
         "EastwatchLE",
         "Ephemeron",
@@ -143,12 +145,15 @@ def main():
         "FractureLE",
         "FrostLE",
         "GlitteringAshesAIE",
+        "GoldenauraAIE",
         "GoldenWall506",
         "GoldenWallLE",
+        "GresvanAIE",
         "HardwireAIE",
         "HonorgroundsLE",
         "IceandChrome506",
         "IceandChromeLE",
+        "InfestationStationAIE",
         "InsideAndOutAIE",
         "InterloperLE",
         "JagannathaAIE",
@@ -175,6 +180,7 @@ def main():
         "RedshiftLE",
         "Reminiscence",
         "RomanticideAIE",
+        "RoyalBloodAIE",
         "Sanglune",
         "SequencerLE",
         "SimulacrumLE",

test/run_example_bots_vs_computer.py

@@ -18,7 +18,7 @@
 from sc2.player import Bot, Computer
 
 # Time limit given in seconds of total in game time
-game_time_limit_vs_computer = 120
+game_time_limit_vs_computer = 240
 
 bot_infos = [
     # Protoss

test/test_pickled_data.py

@@ -13,6 +13,7 @@
 import pickle
 import random
 import sys
+import unittest
 from contextlib import suppress
 from pathlib import Path
 from typing import Any, List, Tuple
@@ -24,6 +25,7 @@
 
 from sc2.bot_ai import BotAI
 from sc2.client import Client
+from sc2.constants import ALL_GAS, CREATION_ABILITY_FIX
 from sc2.data import CloakState, Race
 from sc2.game_data import AbilityData, Cost, GameData
 from sc2.game_info import GameInfo
@@ -919,6 +921,61 @@ def test_units():
     assert scvs.by_tag(scvs[0].tag)
 
 
+def test_exact_creation_ability():
+    try:
+        from sc2.dicts.unit_abilities import UNIT_ABILITIES
+        from sc2.dicts.unit_unit_alias import UNIT_UNIT_ALIAS
+    except ImportError:
+        logger.info(f"Import error: dict sc2/dicts/ are missing!")
+        return
+    test_case = unittest.TestCase()
+    bot: BotAI = get_map_specific_bot(random.choice(MAPS))
+
+    ignore_types = {
+        UnitTypeId.ADEPTPHASESHIFT,
+        UnitTypeId.ARBITERMP,
+        UnitTypeId.BROODLING,
+        UnitTypeId.BYPASSARMORDRONE,
+        UnitTypeId.CORSAIRMP,
+        UnitTypeId.EGG,
+        UnitTypeId.ELSECARO_COLONIST_HUT,
+        UnitTypeId.HERC,
+        UnitTypeId.HERCPLACEMENT,
+        UnitTypeId.INFESTEDTERRANSEGG,
+        UnitTypeId.LARVA,
+        UnitTypeId.NYDUSCANALCREEPER,
+        UnitTypeId.QUEENMP,
+        UnitTypeId.RAVENREPAIRDRONE,
+        UnitTypeId.REPLICANT,
+        UnitTypeId.SCOURGEMP,
+        UnitTypeId.SCOUTMP,
+        UnitTypeId.WARHOUND,
+    }
+
+    unit_types = list(UNIT_UNIT_ALIAS) + list(UNIT_UNIT_ALIAS.values()) + list(UNIT_ABILITIES) + list(ALL_GAS)
+    unit_types_unique_sorted = sorted(set(t.name for t in unit_types))
+    for unit_type_name in unit_types_unique_sorted:
+        unit_type = UnitTypeId[unit_type_name]
+        if unit_type in ignore_types:
+            continue
+
+        if unit_type in [
+            UnitTypeId.ARCHON,
+            UnitTypeId.ASSIMILATORRICH,
+            UnitTypeId.EXTRACTORRICH,
+            UnitTypeId.REFINERYRICH,
+        ]:
+            with test_case.assertRaises(AttributeError):
+                _creation_ability = bot.game_data.units[unit_type.value].creation_ability.exact_id
+            continue
+
+        try:
+            _creation_ability = bot.game_data.units[unit_type.value].creation_ability.exact_id
+        except AttributeError:
+            if unit_type not in CREATION_ABILITY_FIX:
+                assert False, f"Unit type '{unit_type}' missing from CREATION_ABILITY_FIX"
+
+
 def test_dicts():
     # May be missing but that should not fail the tests
     try: