210 lines
6.5 KiB
Python
210 lines
6.5 KiB
Python
# 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)
|