Merge pull request #181 from AresSC2/feat/switch-opening

feat: switch opening

docs/tutorials/build_runner.md

@@ -293,3 +293,15 @@ If omitted the default spawn target is our start location.
 # this is perfectly fine
 - 36 adept
 ```
+
+### Switch openings on the fly
+It's possible to switch openings on the fly, `ares` will attempt to work out
+which build steps have already completed and find a reasonable point in the
+new build order to resume from.
+
+You should pass a valid opening name from your builds yaml file, something like:
+```python
+if self.opponent_is_cheesing:
+    self.build_order_runner.switch_opening("DefensiveOpening")
+```
+Note that if an incorrect opening name is passed here the bot will terminate.

src/ares/behaviors/macro/production_controller.py

@@ -12,7 +12,7 @@
 from sc2.unit import Unit
 from sc2.units import Units
 
-from ares.consts import ADD_ONS, GATEWAY_UNITS, ID, TARGET, TECHLAB_TYPES
+from ares.consts import ADD_ONS, GATEWAY_UNITS, TECHLAB_TYPES
 from ares.dicts.unit_tech_requirement import UNIT_TECH_REQUIREMENT
 
 if TYPE_CHECKING:
@@ -147,7 +147,7 @@ def execute(self, ai: "AresBot", config: dict, mediator: ManagerMediator) -> boo
 
             # we have a worker on route to build this production
             # leave alone for now
-            if self._not_started_but_in_building_tracker(ai, mediator, trained_from):
+            if ai.not_started_but_in_building_tracker(trained_from):
                 continue
 
             # we can afford prod, work out how much prod to support
@@ -290,39 +290,6 @@ def is_flying_production(
                 return True
         return False
 
-    @staticmethod
-    def _not_started_but_in_building_tracker(
-        ai: "AresBot", mediator: ManagerMediator, structure_type: UnitID
-    ) -> bool:
-        """
-        Figures out if worker in on route to build something, and
-        that structure_type doesn't exist yet.
-
-        Parameters
-        ----------
-        ai
-        mediator
-        structure_type
-
-        Returns
-        -------
-
-        """
-        building_tracker: dict = mediator.get_building_tracker_dict
-        for tag, info in building_tracker.items():
-            structure_id: UnitID = building_tracker[tag][ID]
-            if structure_id != structure_type:
-                continue
-
-            target: Point2 = building_tracker[tag][TARGET]
-
-            if not ai.structures.filter(
-                lambda s: cy_distance_to_squared(s.position, target.position) < 1.0
-            ):
-                return True
-
-        return False
-
     def _teching_up(
         self, ai: "AresBot", unit_type_id: UnitID, trained_from: UnitID
     ) -> bool:

src/ares/build_runner/build_order_parser.py

@@ -1,8 +1,9 @@
+from collections import defaultdict
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
 import numpy as np
-from cython_extensions import cy_towards
+from cython_extensions import cy_towards, cy_unit_pending
 from loguru import logger
 from map_analyzer import MapData, Region
 from sc2.data import Race
@@ -29,8 +30,6 @@ class BuildOrderParser:
     -----------
     ai: `AresBot`
         The bot instance.
-    raw_build_order: List[str]
-        The list of build order strings.
     build_order_step_dict: Optional[Dict]
         A dictionary of `BuildOrderStep` objects representing
         the recognized build order commands.
@@ -42,14 +41,15 @@ class BuildOrderParser:
     """
 
     ai: "AresBot"
-    raw_build_order: list[str]
     build_order_step_dict: dict = None
 
     def __post_init__(self) -> None:
         """Initializes the `build_order_step_dict` attribute."""
         self.build_order_step_dict = self._generate_build_step_dict()
 
-    def parse(self) -> list[BuildOrderStep]:
+    def parse(
+        self, raw_build_order: list[str], remove_completed: bool = False
+    ) -> list[BuildOrderStep]:
         """Parses the `raw_build_order` attribute into a list of `BuildOrderStep`.
 
         Returns:
@@ -58,12 +58,15 @@ def parse(self) -> list[BuildOrderStep]:
             The list of `BuildOrderStep` objects parsed from `raw_build_order`.
         """
         build_order: list[BuildOrderStep] = []
-        for raw_step in self.raw_build_order:
+        for raw_step in raw_build_order:
             if isinstance(raw_step, str):
                 build_order = self._parse_string_command(raw_step, build_order)
             elif isinstance(raw_step, dict):
                 build_order = self._parse_dict_command(raw_step, build_order)
 
+        # incase we switched from a different build
+        if remove_completed:
+            build_order = self._remove_completed_steps(build_order)
         return build_order
 
     def _generate_build_step_dict(self) -> dict:
@@ -527,3 +530,79 @@ def _get_target(self, target: Optional[str]) -> Point2:
             case BuildOrderTargetOptions.THIRD:
                 return self.ai.mediator.get_own_expansions[1][0]
         return self.ai.start_location
+
+    def _remove_completed_steps(
+        self, build_order: list[BuildOrderStep]
+    ) -> list[BuildOrderStep]:
+        """
+        Provided a build order, look for steps already completed.
+        This is useful when switching from one opening to another.
+
+        Parameters
+        ----------
+        build_order
+
+        Returns
+        -------
+
+        """
+        indices_to_remove: list[int] = []
+
+        num_same_steps_found: dict[UnitID, int] = defaultdict(int)
+        # pretend we already built things we spawn with
+        # makes working this out easier
+        num_same_steps_found[self.ai.base_townhall_type] = 1
+        num_same_steps_found[UnitID.OVERLORD] = 1
+        num_same_steps_found[self.ai.worker_type] = 12
+
+        for i, step in enumerate(build_order):
+            command: Union[AbilityId, UnitID, UpgradeId] = step.command
+            if command == BuildOrderOptions.WORKER_SCOUT:
+                logger.info(
+                    f"Removing {command} from build order. "
+                    f"Please note worker scouts are always "
+                    f"removed when switching build orders"
+                )
+                indices_to_remove.append(i)
+
+            # remove any steps that chrono the nexus
+            # not ideal but helps build order not getting stuck
+            elif isinstance(command, AbilityId):
+                if (
+                    command == AbilityId.EFFECT_CHRONOBOOST
+                    and step.target == UnitID.NEXUS
+                ):
+                    logger.info(f"Removing {command} from build order")
+                    indices_to_remove.append(i)
+            elif isinstance(command, UnitID):
+                if command in ALL_STRUCTURES:
+                    num_existing: int = len(
+                        self.ai.mediator.get_own_structures_dict[command]
+                    )
+                    on_route: int = int(
+                        self.ai.not_started_but_in_building_tracker(command)
+                    )
+                    total_present: int = num_existing + on_route
+                else:
+                    num_units: int = len(self.ai.mediator.get_own_army_dict[command])
+                    pending: int = cy_unit_pending(self.ai, command)
+                    total_present: int = num_units + pending
+
+                if total_present == 0:
+                    continue
+
+                # while there are less of these steps then what are present
+                if num_same_steps_found[command] < total_present:
+                    logger.info(f"Removing {command} from build order")
+                    num_same_steps_found[command] += 1
+                    indices_to_remove.append(i)
+
+            elif isinstance(command, UpgradeId):
+                if self.ai.pending_or_complete_upgrade(command):
+                    logger.info(f"Removing {command} from build order")
+                    indices_to_remove.append(i)
+
+        for index in sorted(indices_to_remove, reverse=True):
+            del build_order[index]
+
+        return build_order

src/ares/build_runner/build_order_runner.py

@@ -90,50 +90,74 @@ def __init__(
         self.auto_supply_at_supply: int = 200
         self.constant_worker_production_till: int = 0
         self.persistent_worker: bool = True
+
+        self.build_order: list[BuildOrderStep] = []
+        self._build_order_parser: BuildOrderParser = BuildOrderParser(self.ai)
         self._chosen_opening: str = chosen_opening
+        self.configure_opening_from_yml_file(config, chosen_opening)
+
+        self.build_step: int = 0
+        self.current_step_started: bool = False
+        self.current_step_complete: bool = False
+        self._opening_build_completed: bool = False
+        self.current_build_position: Point2 = self.ai.start_location
+        self.assigned_persistent_worker: bool = False
+
+        self._temporary_build_step: int = -1
+
+    def set_build_completed(self) -> None:
+        logger.info("Build order completed")
+        self.mediator.switch_roles(
+            from_role=UnitRole.PERSISTENT_BUILDER, to_role=UnitRole.GATHERING
+        )
+        self._opening_build_completed = True
+
+    def configure_opening_from_yml_file(
+        self, config: dict, opening_name: str, remove_completed: bool = False
+    ) -> None:
         if BUILDS in self.config:
-            build: list[str] = config[BUILDS][chosen_opening][OPENING_BUILD_ORDER]
-            logger.info(f"Running build from yml file: {chosen_opening}")
-            if self.AUTO_SUPPLY_AT_SUPPLY in config[BUILDS][chosen_opening]:
+            assert isinstance(
+                config[BUILDS], dict
+            ), "Opening builds are not configured correctly in the yml file"
+
+            assert opening_name in config[BUILDS].keys(), (
+                f"Trying to parse an opening called {opening_name} but "
+                f"I can't find it. Spelling perhaps?"
+            )
+
+            build: list[str] = config[BUILDS][opening_name][OPENING_BUILD_ORDER]
+            logger.info(
+                f"{self.ai.time_formatted}: Running build from yml file: {opening_name}"
+            )
+            if self.AUTO_SUPPLY_AT_SUPPLY in config[BUILDS][opening_name]:
                 try:
                     self.auto_supply_at_supply = int(
-                        config[BUILDS][chosen_opening][self.AUTO_SUPPLY_AT_SUPPLY]
+                        config[BUILDS][opening_name][self.AUTO_SUPPLY_AT_SUPPLY]
                     )
                 except ValueError as e:
                     logger.warning(f"Error: {e}")
-            if self.CONSTANT_WORKER_PRODUCTION_TILL in config[BUILDS][chosen_opening]:
+            if self.CONSTANT_WORKER_PRODUCTION_TILL in config[BUILDS][opening_name]:
                 try:
                     self.constant_worker_production_till = int(
-                        config[BUILDS][chosen_opening][
+                        config[BUILDS][opening_name][
                             self.CONSTANT_WORKER_PRODUCTION_TILL
                         ]
                     )
                 except ValueError as e:
                     logger.warning(f"Error: {e}")
-            if self.PERSISTENT_WORKER in config[BUILDS][chosen_opening]:
-                self.persistent_worker = config[BUILDS][chosen_opening][
+            if self.PERSISTENT_WORKER in config[BUILDS][opening_name]:
+                self.persistent_worker = config[BUILDS][opening_name][
                     self.PERSISTENT_WORKER
                 ]
-        else:
-            build: list[str] = []
-
-        build_order_parser: BuildOrderParser = BuildOrderParser(self.ai, build)
-        self.build_order: list[BuildOrderStep] = build_order_parser.parse()
-        self.build_step: int = 0
-        self.current_step_started: bool = False
-        self.current_step_complete: bool = False
-        self._opening_build_completed: bool = False
-        self.current_build_position: Point2 = self.ai.start_location
-        self.assigned_persistent_worker: bool = False
 
-        self._temporary_build_step: int = -1
+            self.build_step: int = 0
+            self.current_step_started: bool = False
+            self.current_step_complete: bool = False
+            self.current_build_position: Point2 = self.ai.start_location
+            self.assigned_persistent_worker: bool = False
+            self._temporary_build_step: int = -1
 
-    def set_build_completed(self) -> None:
-        logger.info("Build order completed")
-        self.mediator.switch_roles(
-            from_role=UnitRole.PERSISTENT_BUILDER, to_role=UnitRole.GATHERING
-        )
-        self._opening_build_completed = True
+            self.build_order = self._build_order_parser.parse(build, remove_completed)
 
     def set_step_complete(self, value: UnitID) -> None:
         if (
@@ -147,6 +171,13 @@ def set_step_started(self, value: bool, command) -> None:
         if command == self.build_order[self.build_step].command:
             self.current_step_started = value
 
+    def switch_opening(self, opening_name: str, remove_completed: bool = True) -> None:
+        if self._chosen_opening != opening_name:
+            self._chosen_opening = opening_name
+            self.configure_opening_from_yml_file(
+                self.config, opening_name, remove_completed=remove_completed
+            )
+
     @property
     def build_completed(self) -> bool:
         """

src/ares/custom_bot_ai.py

@@ -17,9 +17,10 @@
 from sc2.unit import Unit
 from sc2.units import Units
 
-from ares.consts import ALL_STRUCTURES
+from ares.consts import ALL_STRUCTURES, ID, TARGET
 from ares.dicts.unit_data import UNIT_DATA
 from ares.dicts.unit_tech_requirement import UNIT_TECH_REQUIREMENT
+from ares.managers.manager_mediator import ManagerMediator
 
 
 class CustomBotAI(BotAI):
@@ -31,6 +32,7 @@ class CustomBotAI(BotAI):
     gas_type: UnitID
     unit_tag_dict: Dict[int, Unit]
     worker_type: UnitID
+    mediator: ManagerMediator
 
     async def on_step(self, iteration: int):  # pragma: no cover
         """Here because all abstract methods have to be implemented.
@@ -108,6 +110,34 @@ def get_total_supply(units: Union[Units, list[Unit]]) -> int:
             ]
         )
 
+    def not_started_but_in_building_tracker(self, structure_type: UnitID) -> bool:
+        """
+        Figures out if worker in on route to build something, and
+        that structure_type doesn't exist yet.
+
+        Parameters
+        ----------
+        structure_type
+
+        Returns
+        -------
+
+        """
+        building_tracker: dict = self.mediator.get_building_tracker_dict
+        for tag, info in building_tracker.items():
+            structure_id: UnitID = building_tracker[tag][ID]
+            if structure_id != structure_type:
+                continue
+
+            target: Point2 = building_tracker[tag][TARGET]
+
+            if not self.structures.filter(
+                lambda s: cy_distance_to_squared(s.position, target.position) < 1.0
+            ):
+                return True
+
+        return False
+
     def pending_or_complete_upgrade(self, upgrade_id: UpgradeId) -> bool:
         if upgrade_id in self.state.upgrades:
             return True