SC2: Python 3.11 compatibility (#1821)

Co-authored-by: Salzkorn

worlds/_sc2common/bot/game_data.py

@@ -0,0 +1,209 @@
+# pylint: disable=W0212
+from __future__ import annotations
+
+from bisect import bisect_left
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Dict, List, Optional, Union
+
+from .data import Attribute, Race
+
+# Set of parts of names of abilities that have no cost
+# E.g every ability that has 'Hold' in its name is free
+FREE_ABILITIES = {"Lower", "Raise", "Land", "Lift", "Hold", "Harvest"}
+
+
+class GameData:
+
+    def __init__(self, data):
+        """
+        :param data:
+        """
+        self.abilities: Dict[int, AbilityData] = {}
+        self.units: Dict[int, UnitTypeData] = {u.unit_id: UnitTypeData(self, u) for u in data.units if u.available}
+        self.upgrades: Dict[int, UpgradeData] = {u.upgrade_id: UpgradeData(self, u) for u in data.upgrades}
+        # Cached UnitTypeIds so that conversion does not take long. This needs to be moved elsewhere if a new GameData object is created multiple times per game
+
+
+class AbilityData:
+
+    @classmethod
+    def id_exists(cls, ability_id):
+        assert isinstance(ability_id, int), f"Wrong type: {ability_id} is not int"
+        if ability_id == 0:
+            return False
+        i = bisect_left(cls.ability_ids, ability_id)  # quick binary search
+        return i != len(cls.ability_ids) and cls.ability_ids[i] == ability_id
+
+    def __init__(self, game_data, proto):
+        self._game_data = game_data
+        self._proto = proto
+
+        # What happens if we comment this out? Should this not be commented out? What is its purpose?
+        assert self.id != 0
+
+    def __repr__(self) -> str:
+        return f"AbilityData(name={self._proto.button_name})"
+
+    @property
+    def link_name(self) -> str:
+        """ For Stimpack this returns 'BarracksTechLabResearch' """
+        return self._proto.link_name
+
+    @property
+    def button_name(self) -> str:
+        """ For Stimpack this returns 'Stimpack' """
+        return self._proto.button_name
+
+    @property
+    def friendly_name(self) -> str:
+        """ For Stimpack this returns 'Research Stimpack' """
+        return self._proto.friendly_name
+
+    @property
+    def is_free_morph(self) -> bool:
+        return any(free in self._proto.link_name for free in FREE_ABILITIES)
+
+    @property
+    def cost(self) -> Cost:
+        return self._game_data.calculate_ability_cost(self.id)
+
+
+class UnitTypeData:
+
+    def __init__(self, game_data: GameData, proto):
+        """
+        :param game_data:
+        :param proto:
+        """
+        self._game_data = game_data
+        self._proto = proto
+
+    def __repr__(self) -> str:
+        return f"UnitTypeData(name={self.name})"
+
+    @property
+    def name(self) -> str:
+        return self._proto.name
+
+    @property
+    def creation_ability(self) -> Optional[AbilityData]:
+        if self._proto.ability_id == 0:
+            return None
+        if self._proto.ability_id not in self._game_data.abilities:
+            return None
+        return self._game_data.abilities[self._proto.ability_id]
+
+    @property
+    def footprint_radius(self) -> Optional[float]:
+        """ See unit.py footprint_radius """
+        if self.creation_ability is None:
+            return None
+        return self.creation_ability._proto.footprint_radius
+
+    @property
+    def attributes(self) -> List[Attribute]:
+        return self._proto.attributes
+
+    def has_attribute(self, attr) -> bool:
+        assert isinstance(attr, Attribute)
+        return attr in self.attributes
+
+    @property
+    def has_minerals(self) -> bool:
+        return self._proto.has_minerals
+
+    @property
+    def has_vespene(self) -> bool:
+        return self._proto.has_vespene
+
+    @property
+    def cargo_size(self) -> int:
+        """ How much cargo this unit uses up in cargo_space """
+        return self._proto.cargo_size
+
+    @property
+    def race(self) -> Race:
+        return Race(self._proto.race)
+
+    @property
+    def cost(self) -> Cost:
+        return Cost(self._proto.mineral_cost, self._proto.vespene_cost, self._proto.build_time)
+
+    @property
+    def cost_zerg_corrected(self) -> Cost:
+        """ This returns 25 for extractor and 200 for spawning pool instead of 75 and 250 respectively """
+        if self.race == Race.Zerg and Attribute.Structure.value in self.attributes:
+            return Cost(self._proto.mineral_cost - 50, self._proto.vespene_cost, self._proto.build_time)
+        return self.cost
+
+
+class UpgradeData:
+
+    def __init__(self, game_data: GameData, proto):
+        """
+        :param game_data:
+        :param proto:
+        """
+        self._game_data = game_data
+        self._proto = proto
+
+    def __repr__(self):
+        return f"UpgradeData({self.name} - research ability: {self.research_ability}, {self.cost})"
+
+    @property
+    def name(self) -> str:
+        return self._proto.name
+
+    @property
+    def research_ability(self) -> Optional[AbilityData]:
+        if self._proto.ability_id == 0:
+            return None
+        if self._proto.ability_id not in self._game_data.abilities:
+            return None
+        return self._game_data.abilities[self._proto.ability_id]
+
+    @property
+    def cost(self) -> Cost:
+        return Cost(self._proto.mineral_cost, self._proto.vespene_cost, self._proto.research_time)
+
+
+@dataclass
+class Cost:
+    """
+    The cost of an action, a structure, a unit or a research upgrade.
+    The time is given in frames (22.4 frames per game second).
+    """
+    minerals: int
+    vespene: int
+    time: Optional[float] = None
+
+    def __repr__(self) -> str:
+        return f"Cost({self.minerals}, {self.vespene})"
+
+    def __eq__(self, other: Cost) -> bool:
+        return self.minerals == other.minerals and self.vespene == other.vespene
+
+    def __ne__(self, other: Cost) -> bool:
+        return self.minerals != other.minerals or self.vespene != other.vespene
+
+    def __bool__(self) -> bool:
+        return self.minerals != 0 or self.vespene != 0
+
+    def __add__(self, other) -> Cost:
+        if not other:
+            return self
+        if not self:
+            return other
+        time = (self.time or 0) + (other.time or 0)
+        return Cost(self.minerals + other.minerals, self.vespene + other.vespene, time=time)
+
+    def __sub__(self, other: Cost) -> Cost:
+        time = (self.time or 0) + (other.time or 0)
+        return Cost(self.minerals - other.minerals, self.vespene - other.vespene, time=time)
+
+    def __mul__(self, other: int) -> Cost:
+        return Cost(self.minerals * other, self.vespene * other, time=self.time)
+
+    def __rmul__(self, other: int) -> Cost:
+        return Cost(self.minerals * other, self.vespene * other, time=self.time)