diff --git a/sc2/bot_ai.py b/sc2/bot_ai.py index f7567d529..250d4f73d 100644 --- a/sc2/bot_ai.py +++ b/sc2/bot_ai.py @@ -5,8 +5,6 @@ from collections import Counter from typing import Any, Dict, List, Optional, Set, Tuple, Union # mypy type checking -from s2clientprotocol import common_pb2 as common_pb - from .cache import property_cache_forever, property_cache_once_per_frame from .data import ActionResult, Alert, Race, Result, Target, race_gas, race_townhalls, race_worker from .game_data import AbilityData, GameData @@ -130,7 +128,7 @@ def main_base_ramp(self) -> "Ramp": Look in game_info.py for more information """ if hasattr(self, "cached_main_base_ramp"): return self.cached_main_base_ramp - # The reason for len(ramp.upper) in {2, 5} is: + # The reason for len(ramp.upper) in {2, 5} is: # ParaSite map has 5 upper points, and most other maps have 2 upper points at the main ramp. # The map Acolyte has 4 upper points at the wrong ramp (which is closest to the start position). try: @@ -157,65 +155,58 @@ def expansion_locations(self) -> Dict[Point2, Units]: # any resource in a group is closer than 6 to any resource of another group # Distance we group resources by - from .helpers.devtools import time_this - - with time_this("expo locations"): - RESOURCE_SPREAD_THRESHOLD = 8.5 - minerals = self.state.mineral_field - geysers = self.state.vespene_geyser - all_resources = minerals | geysers - # Presort resources to get faster clustering - # all_resources.sort(key=lambda resource: resource.position.x ** 2 + resource.position.y ** 2) - all_resources.sort(key=lambda resource: resource.position.x ** 2 + resource.position.y ** 2) - # Create a group for every resource - resource_groups = [[resource] for resource in all_resources] - # Loop the merging process as long as we change something - found_something = True - while found_something: - found_something = False - # Check every combination of two groups - for group_a, group_b in itertools.combinations(resource_groups, 2): - # Check if any pair of resource of these groups is closer than threshold together - if any( - resource_a.distance_to(resource_b) <= RESOURCE_SPREAD_THRESHOLD - for resource_a, resource_b in itertools.product(group_a, group_b) - ): - # Remove the single groups and add the merged group - resource_groups.remove(group_a) - resource_groups.remove(group_b) - resource_groups.append(group_a + group_b) - found_something = True - break - # Distance offsets we apply to center of each resource group to find expansion position - offset_range = 7 - offsets = [(x, y) for x, y in itertools.product(range(-offset_range, offset_range + 1), repeat=2)] - # Dict we want to return - centers = {} - # For every resource group: - for resources in resource_groups: - # Possible expansion points - amount = len(resources) - # Calculate center, round and add 0.5 because expansion location will have (x.5, y.5) - # coordinates because bases have size 5. - center_x = round(sum(resource.position.x for resource in resources) / amount) + 0.5 - center_y = round(sum(resource.position.y for resource in resources) / amount) + 0.5 - possible_points = (Point2((offset[0] + center_x, offset[1] + center_y)) for offset in offsets) - # Filter out points that are too near - possible_points = ( - point - for point in possible_points - # Check if point can be built on - if self._game_info.placement_grid[point.rounded] == 1 - # Check if all resources have enough space to point - and all(point.distance_to(resource) > (7 if resource in geysers else 6) for resource in resources) - ) - # Choose best fitting point - # TODO can we improve this by calculating the distance only one time? - result = min( - possible_points, key=lambda point: sum(point.distance_to(resource) for resource in resources) - ) - centers[result] = resources - return centers + RESOURCE_SPREAD_THRESHOLD = 8.5 + geysers = self.state.vespene_geyser + # Create a group for every resource + resource_groups = [[resource] for resource in self.state.resources] + # Loop the merging process as long as we change something + found_something = True + while found_something: + found_something = False + # Check every combination of two groups + for group_a, group_b in itertools.combinations(resource_groups, 2): + # Check if any pair of resource of these groups is closer than threshold together + if any( + resource_a.distance_to(resource_b) <= RESOURCE_SPREAD_THRESHOLD + for resource_a, resource_b in itertools.product(group_a, group_b) + ): + # Remove the single groups and add the merged group + resource_groups.remove(group_a) + resource_groups.remove(group_b) + resource_groups.append(group_a + group_b) + found_something = True + break + # Distance offsets we apply to center of each resource group to find expansion position + offset_range = 7 + offsets = [ + (x, y) + for x, y in itertools.product(range(-offset_range, offset_range + 1), repeat=2) + if math.hypot(x, y) <= 8 + ] + # Dict we want to return + centers = {} + # For every resource group: + for resources in resource_groups: + # Possible expansion points + amount = len(resources) + # Calculate center, round and add 0.5 because expansion location will have (x.5, y.5) + # coordinates because bases have size 5. + center_x = int(sum(resource.position.x for resource in resources) / amount) + 0.5 + center_y = int(sum(resource.position.y for resource in resources) / amount) + 0.5 + possible_points = (Point2((offset[0] + center_x, offset[1] + center_y)) for offset in offsets) + # Filter out points that are too near + possible_points = ( + point + for point in possible_points + # Check if point can be built on + if self._game_info.placement_grid[point.rounded] == 1 + # Check if all resources have enough space to point + and all(point.distance_to(resource) > (7 if resource in geysers else 6) for resource in resources) + ) + # Choose best fitting point + result = min(possible_points, key=lambda point: sum(point.distance_to(resource) for resource in resources)) + centers[result] = resources + return centers def _correct_zerg_supply(self): """ The client incorrectly rounds zerg supply down instead of up (see diff --git a/sc2/client.py b/sc2/client.py index 039de3e0f..3c269c1e9 100644 --- a/sc2/client.py +++ b/sc2/client.py @@ -138,6 +138,26 @@ async def get_game_data(self) -> GameData: ) return GameData(result.data) + async def dump_data(self, ability_id=True, unit_type_id=True, upgrade_id=True, buff_id=True, effect_id=True): + """ + Dump the game data files + choose what data to dump in the keywords + this function writes to a text file + call it one time in on_step with: + await self._client.dump_data() + """ + result = await self._execute( + data=sc_pb.RequestData( + ability_id=ability_id, + unit_type_id=unit_type_id, + upgrade_id=upgrade_id, + buff_id=buff_id, + effect_id=effect_id, + ) + ) + with open("data_dump.txt", "a") as file: + file.write(str(result.data)) + async def get_game_info(self) -> GameInfo: result = await self._execute(game_info=sc_pb.RequestGameInfo()) return GameInfo(result.game_info) diff --git a/sc2/unit.py b/sc2/unit.py index de4559774..213048a5f 100644 --- a/sc2/unit.py +++ b/sc2/unit.py @@ -43,9 +43,7 @@ def __repr__(self) -> str: return f"UnitOrder({self.ability}, {self.target}, {self.progress})" -class PassengerUnit: - """ Is inherited by the Unit class. Everything in here is also available in Unit. """ - +class Unit: def __init__(self, proto_data): self._proto = proto_data self.cache = {} @@ -302,11 +300,6 @@ def energy_percentage(self) -> Union[int, float]: return 0 return self._proto.energy / self._proto.energy_max - -class Unit(PassengerUnit): - - # All type data is in PassengerUnit. - @property_immutable_cache def is_snapshot(self) -> bool: """ Checks if the unit is only available as a snapshot for the bot. @@ -396,6 +389,11 @@ def is_revealed(self) -> bool: """ Checks if the unit is revealed. """ return self._proto.cloak is CloakState.CloakedDetected.value + @property_immutable_cache + def can_be_attacked(self) -> bool: + """ Checks if the unit is revealed or not cloaked and therefore can be attacked """ + return self._proto.cloak in {CloakState.NotCloaked.value, CloakState.CloakedDetected.value} + @property_immutable_cache def buffs(self) -> Set: """ Returns the set of current buffs the unit has. """ @@ -558,7 +556,8 @@ def is_idle(self) -> bool: return not self.orders def is_using_ability(self, abilities: Union[AbilityId, Set[AbilityId]]) -> bool: - """ Check if the unit is using one of the given abilities. """ + """ Check if the unit is using one of the given abilities. + Only works for own units. """ if not self.orders: return False if isinstance(abilities, AbilityId): @@ -567,12 +566,14 @@ def is_using_ability(self, abilities: Union[AbilityId, Set[AbilityId]]) -> bool: @property_immutable_cache def is_moving(self) -> bool: - """ Checks if the unit is moving. """ + """ Checks if the unit is moving. + Only works for own units. """ return self.is_using_ability(AbilityId.MOVE) @property_immutable_cache def is_attacking(self) -> bool: - """ Checks if the unit is attacking. """ + """ Checks if the unit is attacking. + Only works for own units. """ return self.is_using_ability( { AbilityId.ATTACK, @@ -585,27 +586,32 @@ def is_attacking(self) -> bool: @property_immutable_cache def is_patrolling(self) -> bool: - """ Checks if a unit is patrolling. """ + """ Checks if a unit is patrolling. + Only works for own units. """ return self.is_using_ability(AbilityId.PATROL) @property_immutable_cache def is_gathering(self) -> bool: - """ Checks if a unit is on its way to a mineral field or vespene geyser to mine. """ + """ Checks if a unit is on its way to a mineral field or vespene geyser to mine. + Only works for own units. """ return self.is_using_ability(AbilityId.HARVEST_GATHER) @property_immutable_cache def is_returning(self) -> bool: - """ Checks if a unit is returning from mineral field or vespene geyser to deliver resources to townhall. """ + """ Checks if a unit is returning from mineral field or vespene geyser to deliver resources to townhall. + Only works for own units. """ return self.is_using_ability(AbilityId.HARVEST_RETURN) @property_immutable_cache def is_collecting(self) -> bool: - """ Checks if a unit is gathering or returning. """ + """ Checks if a unit is gathering or returning. + Only works for own units. """ return self.is_using_ability({AbilityId.HARVEST_GATHER, AbilityId.HARVEST_RETURN}) @property_immutable_cache def is_constructing_scv(self) -> bool: - """ Checks if the unit is an SCV that is currently building. """ + """ Checks if the unit is an SCV that is currently building. + Only works for own units. """ return self.is_using_ability( { AbilityId.TERRANBUILD_ARMORY, @@ -626,12 +632,14 @@ def is_constructing_scv(self) -> bool: @property_immutable_cache def is_transforming(self) -> bool: - """ Checks if the unit transforming. """ + """ Checks if the unit transforming. + Only works for own units. """ return self.type_id in transforming and self.is_using_ability(transforming[self.type_id]) @property_immutable_cache def is_repairing(self) -> bool: - """ Checks if the unit is an SCV or MULE that is currently repairing. """ + """ Checks if the unit is an SCV or MULE that is currently repairing. + Only works for own units. """ return self.is_using_ability( {AbilityId.EFFECT_REPAIR, AbilityId.EFFECT_REPAIR_MULE, AbilityId.EFFECT_REPAIR_SCV} ) @@ -653,9 +661,9 @@ def add_on_land_position(self) -> Point2: return self.position.offset(Point2((-2.5, 0.5))) @property_mutable_cache - def passengers(self) -> Set["PassengerUnit"]: + def passengers(self) -> Set["Unit"]: """ Returns the units inside a Bunker, CommandCenter, PlanetaryFortress, Medivac, Nydus, Overlord or WarpPrism. """ - return {PassengerUnit(unit) for unit in self._proto.passengers} + return {Unit(unit) for unit in self._proto.passengers} @property_mutable_cache def passengers_tags(self) -> Set[int]: