diff --git a/game/dcs/aircrafttype.py b/game/dcs/aircrafttype.py index 92e39f1c..7dc7826f 100644 --- a/game/dcs/aircrafttype.py +++ b/game/dcs/aircrafttype.py @@ -48,6 +48,7 @@ from game.utils import ( ) if TYPE_CHECKING: + from game.ato import FlightType from game.missiongenerator.aircraft.flightdata import FlightData from game.missiongenerator.missiondata import MissionData from game.radio.radios import Radio, RadioFrequency, RadioRegistry @@ -198,6 +199,8 @@ class AircraftType(UnitType[Type[FlyingType]]): # when no TGP is mounted on any station. has_built_in_target_pod: bool + task_priorities: dict[FlightType, int] + _by_name: ClassVar[dict[str, AircraftType]] = {} _by_unit_type: ClassVar[dict[type[FlyingType], list[AircraftType]]] = defaultdict( list @@ -322,6 +325,12 @@ class AircraftType(UnitType[Type[FlyingType]]): def iter_props(self) -> Iterator[UnitProperty[Any]]: return UnitProperty.for_aircraft(self.dcs_unit_type) + def capable_of(self, task: FlightType) -> bool: + return task in self.task_priorities + + def task_priority(self, task: FlightType) -> int: + return self.task_priorities[task] + def __setstate__(self, state: dict[str, Any]) -> None: # Update any existing models with new data on load. updated = AircraftType.named(state["name"]) @@ -340,6 +349,12 @@ class AircraftType(UnitType[Type[FlyingType]]): cls._load_all() yield from cls._by_unit_type[dcs_unit_type] + @classmethod + def iter_all(cls) -> Iterator[AircraftType]: + if not cls._loaded: + cls._load_all() + yield from cls._by_name.values() + @staticmethod def each_dcs_type() -> Iterator[Type[FlyingType]]: yield from helicopter_map.values() @@ -422,6 +437,12 @@ class AircraftType(UnitType[Type[FlyingType]]): if prop_overrides is not None: cls._set_props_overrides(prop_overrides, aircraft, data_path) + from game.ato.flighttype import FlightType + + task_priorities: dict[FlightType, int] = {} + for task_name, priority in data.get("tasks", {}).items(): + task_priorities[FlightType(task_name)] = priority + for variant in data.get("variants", [aircraft.id]): yield AircraftType( dcs_unit_type=aircraft, @@ -454,4 +475,8 @@ class AircraftType(UnitType[Type[FlyingType]]): cabin_size=data.get("cabin_size", 10 if aircraft.helicopter else 0), can_carry_crates=data.get("can_carry_crates", aircraft.helicopter), has_built_in_target_pod=data.get("has_built_in_target_pod", False), + task_priorities=task_priorities, ) + + def __hash__(self) -> int: + return hash(self.name)