async def on_start(self):
file_path = self.get_pickle_file_path()
print(f"Saving file to {self.map_name}.xz")
await self.store_data_to_file(file_path)
# Make map visible
await self.client.debug_show_map()
await self.client.debug_control_enemy()
await self.client.debug_god()
# Spawn one of each unit
# await self.client.debug_create_unit([[unit_id, 1, self.game_info.map_center, 1] for unit_id in self.game_data.units])
valid_units: Set[UnitTypeId] = {
UnitTypeId(unit_id)
for unit_id, data in self.game_data.units.items()
if data._proto.race != Race.NoRace and data._proto.race != Race.Random and data._proto.available
# Dont cloak units
and UnitTypeId(unit_id) != UnitTypeId.MOTHERSHIP
and (data._proto.mineral_cost or data._proto.movement_speed or data._proto.weapons)
}
# Create units for self
await self.client.debug_create_unit([[valid_unit, 1, self.start_location, 1] for valid_unit in valid_units])
# Create units for enemy
await self.client.debug_create_unit(
[[valid_unit, 1, self.enemy_start_locations[0], 2] for valid_unit in valid_units]
)
await self._advance_steps(2)
file_path = self.get_combat_file_path()
await self.store_data_to_file(file_path)
await self._client.leave()
return
def tech_requirement(self) -> Optional[UnitTypeId]:
""" Tech-building requirement of buildings - may work for units but unreliably """
if self._proto.tech_requirement == 0:
return None
if self._proto.tech_requirement not in self._game_data.units:
return None
return UnitTypeId(self._proto.tech_requirement)
async def production_executioner(self, actions):
'''
Actions is a vector with the nodes corresponding to:
- for each unit
- priority
They result in the following abilities:
- self.do(self.units(LARVA).random.train(<Unit_type_id>))
'''
mask = actions > 0.5
action_asort = actions.argsort()
action_asort_filtered = action_asort[mask]
action_asort_filtered = np.flip(action_asort_filtered)
for action_idx in action_asort_filtered:
unit_id = C.ZERG_UNITS_LARVATRAINABLE_IDS[action_idx]
unit_type_id = UnitTypeId(unit_id)
ability_id = self._game_data.units[
unit_type_id.value].creation_ability.id
larvas = self.units(UnitTypeId.LARVA)
if larvas.exists:
larva = larvas.random
abilities = await self.get_available_abilities(larva)
if ability_id in abilities:
await self.do(larva.train(unit_type_id))
self.log_action(unit_type_id)
def unit_alias(self) -> Optional[UnitTypeId]:
""" Building type equality, e.g. FlyingOrbitalCommand is the same as OrbitalCommand """
if self._proto.unit_alias == 0:
return None
if self._proto.unit_alias not in self._game_data.units:
return None
""" For flying OrbitalCommand, this returns UnitTypeId.OrbitalCommand """
return UnitTypeId(self._proto.unit_alias)
def generate_unit_alias_dict(data: dict):
ability_data = data["Ability"]
unit_data = data["Unit"]
upgrade_data = data["Upgrade"]
# Load pickled game data files from one of the test files
pickled_file_path = Path(
__file__).parent / "test" / "pickle_data" / "AcropolisLE.xz"
assert pickled_file_path.is_file(
), f"Could not find pickled data file {pickled_file_path}"
logger.info(f"Loading pickled game data file {pickled_file_path}")
with lzma.open(pickled_file_path.absolute(), "rb") as f:
raw_game_data, raw_game_info, raw_observation = pickle.load(f)
game_data = GameData(raw_game_data.data)
all_unit_aliases: Dict[UnitTypeId, UnitTypeId] = OrderedDict2()
all_tech_aliases: Dict[UnitTypeId, Set[UnitTypeId]] = OrderedDict2()
entry: dict
for entry in unit_data:
unit_type_value = entry["id"]
unit_type = UnitTypeId(entry["id"])
current_unit_tech_aliases: Set[UnitTypeId] = OrderedSet2()
assert (
unit_type_value in game_data.units
), f"Unit {unit_type} not listed in game_data.units - perhaps pickled file {pickled_file_path} is outdated?"
unit_alias: int = game_data.units[unit_type_value]._proto.unit_alias
if unit_alias:
# Might be 0 if it has no alias
unit_alias_unit_type_id = UnitTypeId(unit_alias)
all_unit_aliases[unit_type] = unit_alias_unit_type_id
tech_aliases: List[int] = game_data.units[
unit_type_value]._proto.tech_alias
for tech_alias in tech_aliases:
# Might be 0 if it has no alias
unit_alias_unit_type_id = UnitTypeId(tech_alias)
current_unit_tech_aliases.add(unit_alias_unit_type_id)
if current_unit_tech_aliases:
all_tech_aliases[unit_type] = current_unit_tech_aliases
return all_unit_aliases, all_tech_aliases
def get_id_from_name(self, name):
for unit in self.unit:
if unit['name'] == name:
return UnitTypeId(unit['id'])
for ability in self.ability:
if ability['name'] == name:
return ability['id']
def generate_unit_alias_dict(data: dict):
ability_data = data["Ability"]
unit_data = data["Unit"]
upgrade_data = data["Upgrade"]
# Load pickled game data files from one of the test files
path = os.path.dirname(__file__)
pickled_files_folder_path = os.path.join(path, "test", "pickle_data")
pickled_files = os.listdir(pickled_files_folder_path)
random_pickled_file = next(f for f in pickled_files if f.endswith(".xz"))
with lzma.open(
os.path.join(pickled_files_folder_path, random_pickled_file),
"rb") as f:
raw_game_data, raw_game_info, raw_observation = pickle.load(f)
game_data = GameData(raw_game_data.data)
all_unit_aliases: Dict[UnitTypeId, UnitTypeId] = OrderedDict2()
all_tech_aliases: Dict[UnitTypeId, Set[UnitTypeId]] = OrderedDict2()
entry: dict
for entry in unit_data:
unit_type_value = entry["id"]
unit_type = UnitTypeId(entry["id"])
current_unit_tech_aliases: Set[UnitTypeId] = OrderedSet2()
assert unit_type_value in game_data.units, f"Unit {unit_type} not listed in game_data.units"
unit_alias: int = game_data.units[unit_type_value]._proto.unit_alias
if unit_alias:
# Might be 0 if it has no alias
unit_alias_unit_type_id = UnitTypeId(unit_alias)
all_unit_aliases[unit_type] = unit_alias_unit_type_id
tech_aliases: List[int] = game_data.units[
unit_type_value]._proto.tech_alias
for tech_alias in tech_aliases:
# Might be 0 if it has no alias
unit_alias_unit_type_id = UnitTypeId(tech_alias)
current_unit_tech_aliases.add(unit_alias_unit_type_id)
if current_unit_tech_aliases:
all_tech_aliases[unit_type] = current_unit_tech_aliases
return all_unit_aliases, all_tech_aliases
def tech_alias(self) -> Optional[List[UnitTypeId]]:
"""Building tech equality, e.g. OrbitalCommand is the same as CommandCenter
Building tech equality, e.g. Hive is the same as Lair and Hatchery
For Hive, this returns [UnitTypeId.Hatchery, UnitTypeId.Lair]
For SCV, this returns None"""
return_list = [
UnitTypeId(tech_alias) for tech_alias in self._proto.tech_alias if tech_alias in self._game_data.units
]
return return_list if return_list else None
def units_including_pending():
view = np.zeros(len(C.ZERG_UNITS_LARVATRAINABLE_IDS))
for unit in self.units.filter(lambda u: not u.is_structure):
if unit.type_id in C.ZERG_UNITS_LARVATRAINABLE_IDS:
view[C.get_zerg_unit_id(known_enemy_unit)] += 0.01
for unit_id_idx, unit_type_idx in enumerate(
C.ZERG_UNITS_LARVATRAINABLE_IDS):
view[unit_id_idx] += self.already_pending(
UnitTypeId(unit_type_idx)) / 100
return view #softmax ?
def ability_specialization_allowed_for(self, a_id, u_id):
"""Can this unit use this specialization, e.g. do not allow LIFT_COMMANDCENTER for BARRACKS."""
assert isinstance(a_id, int)
assert isinstance(u_id, int)
a_name = AbilityId(a_id).name.upper()
u_name = UnitTypeId(u_id).name
if not self.is_specialization(a_name):
return True
for postfix in ["FLYING", "BURROWED", "MP"]:
u_name = remove_postfix(u_name, postfix)
return u_name in a_name
def get_unit_abilities(data: dict):
ability_data = data["Ability"]
unit_data = data["Unit"]
upgrade_data = data["Upgrade"]
all_unit_abilities: Dict[UnitTypeId, Set[AbilityId]] = OrderedDict2()
entry: dict
for entry in unit_data:
entry_unit_abilities = entry.get("abilities", [])
unit_type = UnitTypeId(entry["id"])
current_collected_unit_abilities: Set[AbilityId] = OrderedSet2()
for ability_info in entry_unit_abilities:
ability_id_value: int = ability_info.get("ability", 0)
if ability_id_value:
ability_id: AbilityId = AbilityId(ability_id_value)
current_collected_unit_abilities.add(ability_id)
# logger.info(unit_type, current_unit_abilities)
if current_collected_unit_abilities:
all_unit_abilities[unit_type] = current_collected_unit_abilities
return all_unit_abilities
def get_trainBuilding(self, ability):
for unit in self.unit:
for abbility in unit['abilities']:
if abbility['ability'] == self.get_trainAbillity(ability):
return UnitTypeId(unit['id'])
ZERG_WORKERS = list(filter(lambda u: u['is_worker'], ZERG_UNITS))
ZERG_MILITARY = list(filter(lambda u: not (u['is_structure'] or u['is_worker']), ZERG_UNITS))
#### ABILITIES
ZERG_ABILITIES_IDS = list(reduce(
lambda acc, val: acc.update(map(lambda a: a['ability'], val['abilities'])) or acc,
ZERG_UNITS,
set()))
ZERG_MILITARY_ABILITIES_IDS = list(reduce(
lambda acc, val: acc.update(map(lambda a: a['ability'], val['abilities'])) or acc,
ZERG_MILITARY,
set()))
ZERG_MILITARY_ABILITIES = list(map(
lambda a_id: next(filter(lambda a: a['id'] == a_id, RAW['Ability'])),
ZERG_MILITARY_ABILITIES_IDS))
ZERG_LARVATRAIN_ABILITIES = list(filter(lambda a: a['name'].startswith('LARVATRAIN_'), RAW['Ability']))
ZERG_ZERGBUILD_ABILITIES = list(filter(lambda a: a['name'].startswith('ZERGBUILD_') and a['name'] != 'ZERGBUILD_CANCEL' and a['name'] != 'ZERGBUILD_CREEPTUMOR', RAW['Ability']))
#### UNITS
ZERG_UNITS_IDS = list(map(lambda u: u['id'], ZERG_UNITS))
ZERG_MILITARY_IDS = list(map(lambda u: u['id'], ZERG_MILITARY))
ZERG_MILITARY_UNIT_TYPEID_IDS = list(map(lambda uid: UnitTypeId(uid), ZERG_MILITARY_IDS))
ZERG_UNITS_LARVATRAINABLE_IDS = list(map(lambda a: a['target']['Morph']['produces'], ZERG_LARVATRAIN_ABILITIES))
ZERG_UNITS_LARVATRAINABLE = list(map(
lambda u_id: next(filter(lambda u: u['id'] == u_id, ZERG_UNITS)),
ZERG_UNITS_LARVATRAINABLE_IDS))
#### BUILDINGS
ZERG_BUILDINGS = list(filter(lambda u: u['is_structure'], ZERG_UNITS))
ZERG_BUILDINGS_ZERGBUILD_IDS = list(map(lambda a: a['target'].popitem()[1]['produces'], ZERG_ZERGBUILD_ABILITIES))
ZERG_BUILDINGS_ZERGBUILD_TYPEIDS = list(map(lambda uid: UnitTypeId(uid), ZERG_BUILDINGS_ZERGBUILD_IDS))
ZERG_BUILDINGS_IDS = list(map(lambda u: u['id'], ZERG_BUILDINGS))
ZERG_BUILDINGS_UNIT_TYPEID_IDS = list(map(lambda uid: UnitTypeId(uid), ZERG_BUILDINGS_IDS))
async def state_step(self):
self.wait_steps -= 1
if self.wait_steps > 0:
return
if self.__state == "Empty":
if len(self.unit_queue) == 0:
with (TARGET_DIR / "upgrade.toml").open("w") as f:
f.write(toml.dumps({"Upgrade": self.data_upgrades}))
with (TARGET_DIR / "unit.toml").open("w") as f:
f.write(toml.dumps({"Unit": self.data_units}))
with (TARGET_DIR / "ability.toml").open("w") as f:
f.write(toml.dumps({"Ability": self.data_abilities}))
await self._client.leave()
return
print("Units left:", len(self.unit_queue))
self.current_unit = self.unit_queue.pop()
await self._client.debug_create_unit(
[[UnitTypeId(self.current_unit), 1, self._game_info.map_center, 1]]
)
self.time_left = 10
self.__state = "WaitCreate"
elif self.__state == "WaitCreate":
if len(self.units) == 0 and self.current_unit == UnitTypeId.LARVA.value:
# Larva cannot be created without a hatchery
await self._client.debug_create_unit(
[[UnitTypeId.HATCHERY, 1, self._game_info.map_center, 1]]
)
self.wait_steps = 10
return
elif len(self.units) == 0:
self.time_left -= 1
if self.time_left < 0:
index = [
i
for i, u in enumerate(self.data_units)
if u["id"] == self.current_unit
][0]
del self.data_units[index]
self.__state = "Clear"
else:
cands = [
u for u in self.units if u._proto.unit_type == self.current_unit
]
if len(cands) == 0:
# Check for some allowed specialization
su = self.units.first.name.upper()
lu = UnitTypeId(self.current_unit).name.upper()
if len(self.units) == 1 and (
su in lu or all(n.startswith("CREEPTUMOR") for n in (su, lu))
):
unit = self.units.first
else:
assert (
False
), f"Invalid self.units (looking for {UnitTypeId(self.current_unit) !r}): {self.units}"
else:
unit = cands[0]
assert unit.is_ready
index = [
i
for i, u in enumerate(self.data_units)
if u["id"] == self.current_unit
][0]
if self.current_unit in [
UnitTypeId.CREEPTUMOR.value,
UnitTypeId.CREEPTUMORQUEEN.value,
]:
# TODO: Handle this properly
# Creep tumors automatically burrow when complete
# CREEPTUMORBURROWED
pass
elif self.current_unit == UnitTypeId.LARVA.value:
# Larva must be selected
unit = self.units(UnitTypeId.LARVA).first
elif self.current_unit in [
UnitTypeId.BARRACKSTECHLAB.value,
UnitTypeId.BARRACKSREACTOR.value,
UnitTypeId.FACTORYTECHLAB.value,
UnitTypeId.FACTORYREACTOR.value,
UnitTypeId.STARPORTTECHLAB.value,
UnitTypeId.STARPORTREACTOR.value,
]:
# Reactors and tech labs are not really part of the building,
# so to get the abilities an appropriate building must be added.
# Bare Reactor and TechLab have no abilities, so not matching them here.
if self.current_unit in [
UnitTypeId.BARRACKSTECHLAB.value,
UnitTypeId.BARRACKSREACTOR.value,
]:
ut = UnitTypeId.BARRACKS
elif self.current_unit in [
UnitTypeId.FACTORYTECHLAB.value,
UnitTypeId.FACTORYREACTOR.value,
]:
ut = UnitTypeId.FACTORY
elif self.current_unit in [
UnitTypeId.STARPORTTECHLAB.value,
UnitTypeId.STARPORTREACTOR.value,
]:
ut = UnitTypeId.STARPORT
else:
assert False, f"Type? {unit.type_id.name}"
if len(self.units) > 1:
assert len(self.units) == 2 and all(
u.is_ready for u in self.units
)
# Building and addon both created
else:
await self._client.debug_create_unit(
[[ut, 1, self._game_info.map_center, 1]]
)
await self._client.debug_kill_unit([unit.tag])
self.wait_steps = 100
self.__state = "BuildAddOn"
return
elif self.data_units[index]["needs_power"]:
# Build pylon for protoss buildings that need it
if len(self.units) > 1:
assert len(self.units) == 2, f"Units: {self.units}"
assert all(u.is_ready for u in self.units)
assert len(self.state.psionic_matrix.sources) == 1
# Pylon already created
else:
if self.current_unit == UnitTypeId.GATEWAY.value:
# Disable autocast of warpgate morph
await self._client.toggle_autocast(
[unit], AbilityId.MORPH_WARPGATE
)
await self._client.debug_create_unit(
[[UnitTypeId.PYLON, 1, self._game_info.map_center, 1]]
)
self.wait_steps = 200
return
else:
assert (
self.current_unit == unit.type_id.value
), f"{self.current_unit} == {unit.type_id.value} ({unit.type_id})"
self.data_units[index]["cargo_capacity"] = if_nonzero(
unit._proto.cargo_space_max
)
self.data_units[index]["max_health"] = unit._proto.health_max
self.data_units[index]["max_shield"] = if_nonzero(
unit._proto.shield_max
)
self.data_units[index]["detection_range"] = if_nonzero(
unit._proto.detect_range
)
self.data_units[index]["start_energy"] = if_nonzero(
unit._proto.energy, int
)
self.data_units[index]["max_energy"] = if_nonzero(
unit._proto.energy_max
)
self.data_units[index]["radius"] = if_nonzero(unit._proto.radius)
# TODO: "placement_size" for buildings
# Provided power radius
power_sources = self.state.psionic_matrix.sources
if len(power_sources) > 0:
assert len(power_sources) == 1
self.data_units[index]["power_radius"] = power_sources[0].radius
# Unit abilities
try:
abilities = (
await self.get_available_abilities(
[unit], ignore_resource_requirements=True
)
)[0]
# No requirements when all tech is locked
self.data_units[index]["abilities"] = [
{"ability": a.value}
for a in abilities
if self.recognizes_ability(a.value)
and self.ability_specialization_allowed_for(
a.value, unit._proto.unit_type
)
]
# See requirement-depending upgrades with tech
await self._client.debug_tech_tree()
self.__state = "TechCheck"
except ValueError as e:
assert "is not a valid AbilityId" in repr(e), repr(e)
# TODO: maybe skip the unit entirely
self.__state = "Clear"
elif self.__state == "BuildAddOn":
assert len(self.units) == 1, f"? {self.units}"
unit = self.units.first
await self.do(unit.build(UnitTypeId(self.current_unit)))
self.wait_steps = 10
self.__state = "BuildAddOnWait"
elif self.__state == "BuildAddOnWait":
assert len(self.units) == 2, f"? {self.units}"
if all(u.is_ready for u in self.units):
self.__state = "WaitCreate"
elif self.__state == "TechCheck":
possible_units = [
u for u in self.units if u._proto.unit_type == self.current_unit
]
if possible_units:
unit = possible_units[0]
assert unit.is_ready
index = [
i
for i, u in enumerate(self.data_units)
if u["id"] == self.current_unit
][0]
abilities = (
await self.get_available_abilities(
[unit], ignore_resource_requirements=True
)
)[0]
print("#", unit)
for a in abilities:
print(">", a)
if not self.recognizes_ability(a.value):
continue
if not self.ability_specialization_allowed_for(
a.value, unit._proto.unit_type
):
continue
if a.value not in [
a["ability"] for a in self.data_units[index]["abilities"]
]:
self.data_units[index]["abilities"].append(
{"requirements": "???", "ability": a.value}
)
# Switch all tech back off
await self._client.debug_tech_tree()
self.__state = "Clear"
elif self.__state == "WaitCreate":
if len(self.units) == 0:
self.time_left -= 1
if self.time_left < 0:
self.__state = "Clear"
elif self.__state == "WaitEmpty":
if len(self.units) > 0:
self.time_left -= 1
if self.time_left < 0:
assert False, "Clear failed"
else:
# Kill broodlings etc
for u in self.state.units:
await self._client.debug_kill_unit([u.tag])
self.wait_steps = 20
else:
self.__state = "Empty"
if self.__state == "Clear":
for u in self.state.units:
await self._client.debug_kill_unit([u.tag])
self.wait_steps = 20
self.current_unit = None
self.__state = "WaitEmpty"
self.time_left = 10
def main():
unit_count = defaultdict(int)
with open("build_order.json", "r") as f:
raw_order = json.load(f)
race = get_race(raw_order)
unit_count[UnitTypeId.OVERLORD] = 1
if race == Race.Protoss:
worker = UnitTypeId.PROBE
townhall = UnitTypeId.NEXUS
elif race == Race.Terran:
worker = UnitTypeId.SCV
townhall = UnitTypeId.COMMANDCENTER
elif race == Race.Zerg:
unit_count[UnitTypeId.OVERLORD] = 1
worker = UnitTypeId.DRONE
townhall = UnitTypeId.HATCHERY
else:
print("Unknown race")
return
unit_count[worker] = 12
unit_count[townhall] = 1
text_timings = "WarnBuildMacro(["
text = ""
last_line: Optional[str] = None
last_id: Optional[Union[UnitTypeId, UpgradeId]] = None
replacer_dict = action_replacement_dict()
for order in raw_order:
line = ""
if order["type"] in {"unit", "structure", "worker"}:
id = UnitTypeId(order["id"])
unit_count[id] += 1
line += f"BuildId({str(id)}, to_count={unit_count[id]})"
if order["type"] == "structure":
text_timings += f"({str(id)}, {unit_count[id]}, {order['frame'] // 22.4}),"
elif order["type"] == "upgrade":
id = UpgradeId(order["id"])
line += f"Tech({str(id)})"
elif order["type"] == "action" and order["name"] in replacer_dict:
id = None
line += replacer_dict[order["name"]]
else:
id = None
line += "# " + order["name"]
line += ","
if last_line is not None and (last_id is None or last_id != id):
text += last_line + "\n"
last_line = line
last_id = id
text += last_line + "\n"
text_timings += "], []),"
print(text_timings)
print(text)
def get_upgrade_abilities(data):
ability_data = data["Ability"]
unit_data = data["Unit"]
upgrade_data = data["Upgrade"]
ability_to_upgrade_dict: Dict[AbilityId, UpgradeId] = OrderedDict2()
"""
We want to be able to research an upgrade by doing
await self.can_research(UpgradeId, return_idle_structures=True) -> returns list of idle structures that can research it
So we need to assign each upgrade id one building type, and its research ability and requirements (e.g. armory for infantry level 2)
"""
# Collect all upgrades and their corresponding abilities
entry: dict
for entry in ability_data:
if isinstance(entry.get("target", {}), str):
continue
ability_id: AbilityId = AbilityId(entry["id"])
researched_ability_id: UnitTypeId
upgrade_id_value: int = entry.get("target",
{}).get("Research",
{}).get("upgrade", 0)
if upgrade_id_value:
upgrade_id: UpgradeId = UpgradeId(upgrade_id_value)
ability_to_upgrade_dict[ability_id] = upgrade_id
"""
unit_research_abilities = {
UnitTypeId.ENGINEERINGBAY: {
UpgradeId.TERRANINFANTRYWEAPONSLEVEL1:
{
"ability": AbilityId.ENGINEERINGBAYRESEARCH_TERRANINFANTRYWEAPONSLEVEL1,
"required_building": None,
"requires_power": False, # If a pylon nearby is required
},
UpgradeId.TERRANINFANTRYWEAPONSLEVEL2: {
"ability": AbilityId.ENGINEERINGBAYRESEARCH_TERRANINFANTRYWEAPONSLEVEL2,
"required_building": UnitTypeId.ARMORY,
"requires_power": False, # If a pylon nearby is required
},
}
}
"""
unit_research_abilities = OrderedDict2()
for entry in unit_data:
unit_abilities = entry.get("abilities", [])
unit_type = UnitTypeId(entry["id"])
if unit_type == UnitTypeId.TECHLAB:
continue
current_unit_research_abilities = OrderedDict2()
for ability_info in unit_abilities:
ability_id_value: int = ability_info.get("ability", 0)
if ability_id_value:
ability_id: AbilityId = AbilityId(ability_id_value)
# Upgrade is not a known upgrade ability
if ability_id not in ability_to_upgrade_dict:
continue
required_building = None
required_upgrade = None
requirements = ability_info.get("requirements", [])
if requirements:
req_building_id_value = next(
(req["building"]
for req in requirements if req.get("building", 0)),
None)
if req_building_id_value:
req_building_id = UnitTypeId(req_building_id_value)
required_building = req_building_id
req_upgrade_id_value = next(
(req["upgrade"]
for req in requirements if req.get("upgrade", 0)),
None)
if req_upgrade_id_value:
req_upgrade_id = UpgradeId(req_upgrade_id_value)
required_upgrade = req_upgrade_id
requires_power = entry.get("needs_power", False)
resulting_upgrade = ability_to_upgrade_dict[ability_id]
research_info = {"ability": ability_id}
if required_building:
research_info["required_building"] = required_building
if required_upgrade:
research_info["required_upgrade"] = required_upgrade
if requires_power:
research_info["requires_power"] = requires_power
current_unit_research_abilities[
resulting_upgrade] = research_info
if current_unit_research_abilities:
unit_research_abilities[
unit_type] = current_unit_research_abilities
return unit_research_abilities
def id(self) -> UnitTypeId:
return UnitTypeId(self._proto.unit_id)
def get_unit_train_build_abilities(data):
ability_data = data["Ability"]
unit_data = data["Unit"]
upgrade_data = data["Upgrade"]
# From which abilities can a unit be trained
train_abilities: Dict[UnitTypeId, Set[AbilityId]] = OrderedDict2()
# If the ability requires a placement position
ability_requires_placement: Set[AbilityId] = set()
# Map ability to unittypeid
ability_to_unittypeid_dict: Dict[AbilityId, UnitTypeId] = OrderedDict2()
# From which abilities can a unit be morphed
# unit_morph_abilities: Dict[UnitTypeId, Set[AbilityId]] = {}
entry: dict
for entry in ability_data:
"""
"target": "PointOrUnit"
"""
if isinstance(entry.get("target", {}), str):
continue
ability_id: AbilityId = AbilityId(entry["id"])
created_unit_type_id: UnitTypeId
# Check if it is a unit train ability
requires_placement = False
train_unit_type_id_value: int = entry.get("target", {}).get(
"Train", {}).get("produces", 0)
train_place_unit_type_id_value: int = entry.get("target", {}).get(
"TrainPlace", {}).get("produces", 0)
morph_unit_type_id_value: int = entry.get("target", {}).get(
"Morph", {}).get("produces", 0)
build_unit_type_id_value: int = entry.get("target", {}).get(
"Build", {}).get("produces", 0)
build_on_unit_unit_type_id_value: int = entry.get("target", {}).get(
"BuildOnUnit", {}).get("produces", 0)
if not train_unit_type_id_value and train_place_unit_type_id_value:
train_unit_type_id_value = train_place_unit_type_id_value
requires_placement = True
# Collect larva morph abilities, and one way morphs (exclude burrow, hellbat morph, siege tank siege)
# Also doesnt include building addons
if not train_unit_type_id_value and (
"LARVATRAIN_" in ability_id.name or ability_id in {
AbilityId.MORPHTOBROODLORD_BROODLORD,
AbilityId.MORPHZERGLINGTOBANELING_BANELING,
AbilityId.MORPHTORAVAGER_RAVAGER,
AbilityId.MORPH_LURKER,
AbilityId.UPGRADETOLAIR_LAIR,
AbilityId.UPGRADETOHIVE_HIVE,
AbilityId.UPGRADETOGREATERSPIRE_GREATERSPIRE,
AbilityId.UPGRADETOORBITAL_ORBITALCOMMAND,
AbilityId.UPGRADETOPLANETARYFORTRESS_PLANETARYFORTRESS,
AbilityId.MORPH_OVERLORDTRANSPORT,
AbilityId.MORPH_OVERSEER,
}):
# If all morph units are used, unit_trained_from.py will be "wrong" because it will list that a siege tank can be trained from siegetanksieged and similar:
# UnitTypeId.SIEGETANK: {UnitTypeId.SIEGETANKSIEGED, UnitTypeId.FACTORY},
# if not train_unit_type_id_value and morph_unit_type_id_value:
train_unit_type_id_value = morph_unit_type_id_value
# Add all build abilities, like construct buildings and train queen (exception)
if not train_unit_type_id_value and build_unit_type_id_value:
train_unit_type_id_value = build_unit_type_id_value
if "BUILD_" in entry["name"]:
requires_placement = True
# Add build gas building (refinery, assimilator, extractor)
# TODO: target needs to be a unit, not a position, but i dont want to store an extra line just for this - needs to be an exception in bot_ai.py
if not train_unit_type_id_value and build_on_unit_unit_type_id_value:
train_unit_type_id_value = build_on_unit_unit_type_id_value
if train_unit_type_id_value:
created_unit_type_id = UnitTypeId(train_unit_type_id_value)
if created_unit_type_id not in train_abilities:
train_abilities[created_unit_type_id] = {ability_id}
else:
train_abilities[created_unit_type_id].add(ability_id)
if requires_placement:
ability_requires_placement.add(ability_id)
ability_to_unittypeid_dict[ability_id] = created_unit_type_id
"""
unit_train_abilities = {
UnitTypeId.GATEWAY: {
UnitTypeId.ADEPT: {
"ability": AbilityId.TRAIN_ADEPT,
"requires_techlab": False,
"required_building": UnitTypeId.CYBERNETICSCORE, # Or None
"requires_placement_position": False, # True for warp gate
"requires_power": True, # If a pylon nearby is required
},
UnitTypeId.Zealot: {
"ability": AbilityId.GATEWAYTRAIN_ZEALOT,
...
}
}
}
"""
unit_train_abilities: Dict[UnitTypeId,
Dict[str, Union[AbilityId, bool,
UnitTypeId]]] = OrderedDict2()
for entry in unit_data:
unit_abilities = entry.get("abilities", [])
unit_type = UnitTypeId(entry["id"])
current_unit_train_abilities = OrderedDict2()
for ability_info in unit_abilities:
ability_id_value: int = ability_info.get("ability", 0)
if ability_id_value:
ability_id: AbilityId = AbilityId(ability_id_value)
# Ability is not a train ability
if ability_id not in ability_to_unittypeid_dict:
continue
requires_techlab: bool = False
required_building: Optional[UnitTypeId] = None
requires_placement_position: bool = False
requires_power: bool = False
"""
requirements = [
{
"addon": 5
},
{
"building": 29
}
]
"""
requirements: List[Dict[str, int]] = ability_info.get(
"requirements", [])
if requirements:
# Assume train abilities only have one tech building requirement; thors requiring armory and techlab is seperatedly counted
assert (
len([
req
for req in requirements if req.get("building", 0)
]) <= 1
), f"Error: Building {unit_type} has more than one tech requirements with train ability {ability_id}"
# UnitTypeId 5 == Techlab
requires_techlab: bool = any(req for req in requirements
if req.get("addon", 0) == 5)
requires_tech_builing_id_value: int = next(
(req["building"]
for req in requirements if req.get("building", 0)), 0)
if requires_tech_builing_id_value:
required_building = UnitTypeId(
requires_tech_builing_id_value)
if ability_id in ability_requires_placement:
requires_placement_position = True
requires_power = entry.get("needs_power", False)
resulting_unit = ability_to_unittypeid_dict[ability_id]
ability_dict = {"ability": ability_id}
# Only add boolean values and tech requirement if they actually exist, to make the resulting dict file smaller
if requires_techlab:
ability_dict["requires_techlab"] = requires_techlab
if required_building:
ability_dict["required_building"] = required_building
if requires_placement_position:
ability_dict[
"requires_placement_position"] = requires_placement_position
if requires_power:
ability_dict["requires_power"] = requires_power
current_unit_train_abilities[resulting_unit] = ability_dict
if current_unit_train_abilities:
unit_train_abilities[unit_type] = current_unit_train_abilities
return unit_train_abilities