Add squadron selector to flight creator.

Fixes https://github.com/dcs-liberation/dcs_liberation/issues/1142

qt_ui/windows/mission/flight/QFlightCreator.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Optional
+from typing import Optional, Type
 
 from PySide2.QtCore import Qt, Signal
 from PySide2.QtWidgets import (
@@ -11,9 +11,10 @@ from PySide2.QtWidgets import (
     QVBoxLayout,
     QLineEdit,
 )
-from dcs.planes import PlaneType
+from dcs.unittype import FlyingType
 
 from game import Game
+from game.squadrons import Squadron
 from game.theater import ControlPoint, OffMapSpawn
 from gen.ato import Package
 from gen.flights.flight import Flight
@@ -24,6 +25,7 @@ from qt_ui.widgets.combos.QAircraftTypeSelector import QAircraftTypeSelector
 from qt_ui.widgets.combos.QArrivalAirfieldSelector import QArrivalAirfieldSelector
 from qt_ui.widgets.combos.QFlightTypeComboBox import QFlightTypeComboBox
 from qt_ui.widgets.combos.QOriginAirfieldSelector import QOriginAirfieldSelector
+from qt_ui.windows.mission.flight.SquadronSelector import SquadronSelector
 
 
 class QFlightCreator(QDialog):
@@ -56,6 +58,14 @@ class QFlightCreator(QDialog):
         self.aircraft_selector.currentIndexChanged.connect(self.on_aircraft_changed)
         layout.addLayout(QLabeledWidget("Aircraft:", self.aircraft_selector))
 
+        self.squadron_selector = SquadronSelector(
+            self.game.air_wing_for(player=True),
+            self.task_selector.currentData(),
+            self.aircraft_selector.currentData(),
+        )
+        self.squadron_selector.setCurrentIndex(0)
+        layout.addLayout(QLabeledWidget("Squadron:", self.squadron_selector))
+
         self.departure = QOriginAirfieldSelector(
             self.game.aircraft_inventory,
             [cp for cp in game.theater.controlpoints if cp.captured],
@@ -133,13 +143,16 @@ class QFlightCreator(QDialog):
         self.custom_name_text = text
 
     def verify_form(self) -> Optional[str]:
-        aircraft: PlaneType = self.aircraft_selector.currentData()
+        aircraft: Optional[Type[FlyingType]] = self.aircraft_selector.currentData()
-        origin: ControlPoint = self.departure.currentData()
+        squadron: Optional[Squadron] = self.squadron_selector.currentData()
-        arrival: ControlPoint = self.arrival.currentData()
+        origin: Optional[ControlPoint] = self.departure.currentData()
-        divert: ControlPoint = self.divert.currentData()
+        arrival: Optional[ControlPoint] = self.arrival.currentData()
+        divert: Optional[ControlPoint] = self.divert.currentData()
         size: int = self.flight_size_spinner.value()
         if aircraft is None:
             return "You must select an aircraft type."
+        if squadron is None:
+            return "You must select a squadron."
         if not origin.captured:
             return f"{origin.name} is not owned by your coalition."
         if arrival is not None and not arrival.captured:
@@ -164,7 +177,7 @@ class QFlightCreator(QDialog):
             return
 
         task = self.task_selector.currentData()
-        aircraft = self.aircraft_selector.currentData()
+        squadron = self.squadron_selector.currentData()
         origin = self.departure.currentData()
         arrival = self.arrival.currentData()
         divert = self.divert.currentData()
@@ -176,7 +189,7 @@ class QFlightCreator(QDialog):
         flight = Flight(
             self.package,
             self.country,
-            self.game.air_wing_for(player=True).squadron_for(aircraft),
+            squadron,
             size,
             task,
             self.start_type.currentText(),
@@ -200,6 +213,9 @@ class QFlightCreator(QDialog):
 
     def on_aircraft_changed(self, index: int) -> None:
         new_aircraft = self.aircraft_selector.itemData(index)
+        self.squadron_selector.update_items(
+            self.task_selector.currentData(), new_aircraft
+        )
         self.departure.change_aircraft(new_aircraft)
         self.arrival.change_aircraft(new_aircraft)
         self.divert.change_aircraft(new_aircraft)
@@ -223,6 +239,9 @@ class QFlightCreator(QDialog):
             self.task_selector.currentData(),
             self.game.aircraft_inventory.available_types_for_player,
         )
+        self.squadron_selector.update_items(
+            self.task_selector.currentData(), self.aircraft_selector.currentData()
+        )
 
     def update_max_size(self, available: int) -> None:
         self.flight_size_spinner.setMaximum(min(available, 4))

qt_ui/windows/mission/flight/SquadronSelector.py

@@ -0,0 +1,48 @@
+"""Combo box for selecting squadrons."""
+from typing import Type, Optional
+
+from PySide2.QtWidgets import QComboBox
+from dcs.unittype import FlyingType
+
+from game.squadrons import AirWing
+from gen.flights.flight import FlightType
+
+
+class SquadronSelector(QComboBox):
+    """Combo box for selecting squadrons compatible with the given requirements."""
+
+    def __init__(
+        self,
+        air_wing: AirWing,
+        task: Optional[FlightType],
+        aircraft: Optional[Type[FlyingType]],
+    ) -> None:
+        super().__init__()
+        self.air_wing = air_wing
+
+        self.model().sort(0)
+        self.setSizeAdjustPolicy(self.AdjustToContents)
+        self.update_items(task, aircraft)
+
+    def update_items(
+        self, task: Optional[FlightType], aircraft: Optional[Type[FlyingType]]
+    ) -> None:
+        current_squadron = self.currentData()
+        self.clear()
+        if task is None:
+            self.addItem("No task selected", None)
+            return
+        if aircraft is None:
+            self.addItem("No aircraft selected", None)
+            return
+
+        for squadron in self.air_wing.squadrons_for(aircraft):
+            if task in squadron.mission_types:
+                self.addItem(f"{squadron}", squadron)
+
+        if self.count() == 0:
+            self.addItem("No capable aircraft available", None)
+            return
+
+        if current_squadron is not None:
+            self.setCurrentText(f"{current_squadron}")