From 1795ed7617375024b297e05cf5ee9d1be82c1130 Mon Sep 17 00:00:00 2001
From: Dan Albert <dan@gingerhq.net>
Date: Thu, 27 May 2021 20:27:45 -0700
Subject: [PATCH] Limit squadron tasks to those of the aircraft.

https://github.com/dcs-liberation/dcs_liberation/issues/276
---
 game/squadrons.py                   | 17 ++++++++++++++---
 gen/flights/ai_flight_planner_db.py |  8 ++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/game/squadrons.py b/game/squadrons.py
index 6d5bbba5..c73e5bed 100644
--- a/game/squadrons.py
+++ b/game/squadrons.py
@@ -177,6 +177,7 @@ class Squadron:
 
     @classmethod
     def from_yaml(cls, path: Path, game: Game, player: bool) -> Squadron:
+        from gen.flights.ai_flight_planner_db import tasks_for_aircraft
         from gen.flights.flight import FlightType
 
         with path.open() as squadron_file:
@@ -189,6 +190,16 @@ class Squadron:
         pilots = [Pilot(n, player=False) for n in data.get("pilots", [])]
         pilots.extend([Pilot(n, player=True) for n in data.get("players", [])])
 
+        mission_types = [FlightType.from_name(n) for n in data["mission_types"]]
+        tasks = tasks_for_aircraft(unit_type)
+        for mission_type in list(mission_types):
+            if mission_type not in tasks:
+                logging.error(
+                    f"Squadron has mission type {mission_type} but {unit_type} is not "
+                    f"capable of that task: {path}"
+                )
+                mission_types.remove(mission_type)
+
         return Squadron(
             name=data["name"],
             nickname=data["nickname"],
@@ -196,7 +207,7 @@ class Squadron:
             role=data["role"],
             aircraft=unit_type,
             livery=data.get("livery"),
-            mission_types=tuple(FlightType.from_name(n) for n in data["mission_types"]),
+            mission_types=tuple(mission_types),
             pilots=pilots,
             game=game,
             player=player,
@@ -262,7 +273,7 @@ class SquadronLoader:
 
 class AirWing:
     def __init__(self, game: Game, player: bool) -> None:
-        from gen.flights.flight import FlightType
+        from gen.flights.ai_flight_planner_db import tasks_for_aircraft
 
         self.game = game
         self.player = player
@@ -280,7 +291,7 @@ class AirWing:
                     role="Flying Squadron",
                     aircraft=aircraft,
                     livery=None,
-                    mission_types=tuple(FlightType),
+                    mission_types=tuple(tasks_for_aircraft(aircraft)),
                     pilots=[],
                     game=game,
                     player=player,
diff --git a/gen/flights/ai_flight_planner_db.py b/gen/flights/ai_flight_planner_db.py
index f0c94171..9170edf7 100644
--- a/gen/flights/ai_flight_planner_db.py
+++ b/gen/flights/ai_flight_planner_db.py
@@ -426,3 +426,11 @@ def aircraft_for_task(task: FlightType) -> List[Type[FlyingType]]:
     else:
         logging.error(f"Unplannable flight type: {task}")
         return []
+
+
+def tasks_for_aircraft(aircraft: Type[FlyingType]) -> list[FlightType]:
+    tasks = []
+    for task in FlightType:
+        if aircraft in aircraft_for_task(task):
+            tasks.append(task)
+    return tasks