From 2df17c32cd37e8eaa32054eae18b47a7a99b7984 Mon Sep 17 00:00:00 2001
From: Dan Albert <dan@gingerhq.net>
Date: Sat, 22 May 2021 17:05:53 -0700
Subject: [PATCH] Clean up aircraft selector.

---
 qt_ui/widgets/combos/QAircraftTypeSelector.py | 76 +++----------------
 .../windows/mission/flight/QFlightCreator.py  |  2 +-
 2 files changed, 11 insertions(+), 67 deletions(-)

diff --git a/qt_ui/widgets/combos/QAircraftTypeSelector.py b/qt_ui/widgets/combos/QAircraftTypeSelector.py
index e5919b3c..b7949aa2 100644
--- a/qt_ui/widgets/combos/QAircraftTypeSelector.py
+++ b/qt_ui/widgets/combos/QAircraftTypeSelector.py
@@ -2,15 +2,12 @@
 from typing import Iterable, Type
 
 from PySide2.QtWidgets import QComboBox
-
 from dcs.unittype import FlyingType
 
+from game import db
+from gen.flights.ai_flight_planner_db import aircraft_for_task
 from gen.flights.flight import FlightType
 
-import gen.flights.ai_flight_planner_db
-
-from game import Game, db
-
 
 class QAircraftTypeSelector(QComboBox):
     """Combo box for selecting among the given aircraft types."""
@@ -19,77 +16,24 @@ class QAircraftTypeSelector(QComboBox):
         self,
         aircraft_types: Iterable[Type[FlyingType]],
         country: str,
-        mission_type: str,
+        mission_type: FlightType,
     ) -> None:
         super().__init__()
 
         self.model().sort(0)
         self.setSizeAdjustPolicy(self.AdjustToContents)
         self.country = country
-        self.updateItems(mission_type, aircraft_types)
+        self.update_items(mission_type, aircraft_types)
 
-    def updateItems(self, mission_type: str, aircraft_types):
+    def update_items(self, mission_type: FlightType, aircraft_types):
         current_aircraft = self.currentData()
         self.clear()
         for aircraft in aircraft_types:
-            if mission_type in [
-                FlightType.BARCAP,
-                FlightType.ESCORT,
-                FlightType.INTERCEPTION,
-                FlightType.SWEEP,
-                FlightType.TARCAP,
-            ]:
-                if aircraft in gen.flights.ai_flight_planner_db.CAP_CAPABLE:
-                    self.addItem(
-                        f"{db.unit_get_expanded_info(self.country, aircraft, 'name')}",
-                        userData=aircraft,
-                    )
-            elif mission_type in [
-                FlightType.CAS,
-                FlightType.BAI,
-                FlightType.OCA_AIRCRAFT,
-            ]:
-                if aircraft in gen.flights.ai_flight_planner_db.CAS_CAPABLE:
-                    self.addItem(
-                        f"{db.unit_get_expanded_info(self.country, aircraft, 'name')}",
-                        userData=aircraft,
-                    )
-            elif mission_type in [FlightType.SEAD]:
-                if aircraft in gen.flights.ai_flight_planner_db.SEAD_CAPABLE:
-                    self.addItem(
-                        f"{db.unit_get_expanded_info(self.country, aircraft, 'name')}",
-                        userData=aircraft,
-                    )
-            elif mission_type in [FlightType.DEAD]:
-                if aircraft in gen.flights.ai_flight_planner_db.DEAD_CAPABLE:
-                    self.addItem(
-                        f"{db.unit_get_expanded_info(self.country, aircraft, 'name')}",
-                        userData=aircraft,
-                    )
-            elif mission_type in [FlightType.STRIKE]:
-                if aircraft in gen.flights.ai_flight_planner_db.STRIKE_CAPABLE:
-                    self.addItem(
-                        f"{db.unit_get_expanded_info(self.country, aircraft, 'name')}",
-                        userData=aircraft,
-                    )
-            elif mission_type in [FlightType.ANTISHIP]:
-                if aircraft in gen.flights.ai_flight_planner_db.ANTISHIP_CAPABLE:
-                    self.addItem(
-                        f"{db.unit_get_expanded_info(self.country, aircraft, 'name')}",
-                        userData=aircraft,
-                    )
-            elif mission_type in [FlightType.OCA_RUNWAY]:
-                if aircraft in gen.flights.ai_flight_planner_db.RUNWAY_ATTACK_CAPABLE:
-                    self.addItem(
-                        f"{db.unit_get_expanded_info(self.country, aircraft, 'name')}",
-                        userData=aircraft,
-                    )
-            elif mission_type in [FlightType.AEWC]:
-                if aircraft in gen.flights.ai_flight_planner_db.AEWC_CAPABLE:
-                    self.addItem(
-                        f"{db.unit_get_expanded_info(self.country, aircraft, 'name')}",
-                        userData=aircraft,
-                    )
+            if aircraft in aircraft_for_task(mission_type):
+                self.addItem(
+                    f"{db.unit_get_expanded_info(self.country, aircraft, 'name')}",
+                    userData=aircraft,
+                )
         current_aircraft_index = self.findData(current_aircraft)
         if current_aircraft_index != -1:
             self.setCurrentIndex(current_aircraft_index)
diff --git a/qt_ui/windows/mission/flight/QFlightCreator.py b/qt_ui/windows/mission/flight/QFlightCreator.py
index f4aa743c..bf388e36 100644
--- a/qt_ui/windows/mission/flight/QFlightCreator.py
+++ b/qt_ui/windows/mission/flight/QFlightCreator.py
@@ -211,7 +211,7 @@ class QFlightCreator(QDialog):
                 self.restore_start_type = None
 
     def on_task_changed(self) -> None:
-        self.aircraft_selector.updateItems(
+        self.aircraft_selector.update_items(
             self.task_selector.currentData(),
             self.game.aircraft_inventory.available_types_for_player,
         )