Spaces:
Runtime error
Runtime error
File size: 2,508 Bytes
07ed12b dd96d2f 07ed12b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | from __future__ import annotations
from typing import TYPE_CHECKING
from .buildings import BuildingType, BuildingStatus
from .units import UnitType
if TYPE_CHECKING:
from .state import PlayerState
# unit_type -> required active buildings
UNIT_REQUIREMENTS: dict[UnitType, list[BuildingType]] = {
UnitType.SCV: [BuildingType.COMMAND_CENTER],
UnitType.MARINE: [BuildingType.BARRACKS],
UnitType.MEDIC: [BuildingType.BARRACKS, BuildingType.ENGINEERING_BAY],
UnitType.GOLIATH: [BuildingType.FACTORY],
UnitType.TANK: [BuildingType.FACTORY, BuildingType.ARMORY],
UnitType.WRAITH: [BuildingType.STARPORT],
}
# building_type -> required active buildings (to be able to construct it)
BUILDING_REQUIREMENTS: dict[BuildingType, list[BuildingType]] = {
BuildingType.COMMAND_CENTER: [],
BuildingType.SUPPLY_DEPOT: [],
BuildingType.REFINERY: [],
BuildingType.BARRACKS: [BuildingType.SUPPLY_DEPOT],
BuildingType.ENGINEERING_BAY: [BuildingType.BARRACKS],
BuildingType.FACTORY: [BuildingType.BARRACKS],
BuildingType.ARMORY: [BuildingType.FACTORY],
BuildingType.STARPORT: [BuildingType.FACTORY],
}
# unit_type -> building that produces it
PRODUCTION_SOURCES: dict[UnitType, BuildingType] = {
UnitType.SCV: BuildingType.COMMAND_CENTER,
UnitType.MARINE: BuildingType.BARRACKS,
UnitType.MEDIC: BuildingType.BARRACKS,
UnitType.GOLIATH: BuildingType.FACTORY,
UnitType.TANK: BuildingType.FACTORY,
UnitType.WRAITH: BuildingType.STARPORT,
}
def _active(player: "PlayerState", bt: BuildingType) -> bool:
return any(
b.building_type == bt
and b.status not in (BuildingStatus.CONSTRUCTING, BuildingStatus.DESTROYED)
for b in player.buildings.values()
)
def can_build(bt: BuildingType, player: "PlayerState") -> bool:
return all(_active(player, req) for req in BUILDING_REQUIREMENTS[bt])
def can_train(ut: UnitType, player: "PlayerState") -> bool:
return all(_active(player, req) for req in UNIT_REQUIREMENTS[ut])
def missing_for_build(bt: BuildingType, player: "PlayerState") -> list[BuildingType]:
return [req for req in BUILDING_REQUIREMENTS[bt] if not _active(player, req)]
def missing_for_train(ut: UnitType, player: "PlayerState") -> list[BuildingType]:
return [req for req in UNIT_REQUIREMENTS[ut] if not _active(player, req)]
def get_producer(ut: UnitType) -> BuildingType:
return PRODUCTION_SOURCES[ut]
|