Clean up custom loadout interface.

Wraps the pydcs data in a real type so we don't need to spread the
reflection all over.

game/data/weapons.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Dict, Iterator, Set, Tuple, Type, Union, cast
+
+from dcs.unitgroup import FlyingGroup
+from dcs.unittype import FlyingType
+
+
+PydcsWeapon = Dict[str, Union[int, str]]
+PydcsWeaponAssignment = Tuple[int, PydcsWeapon]
+
+
+@dataclass(frozen=True)
+class Weapon:
+    """Wraps a pydcs weapon dict in a hashable type."""
+
+    cls_id: str
+    name: str
+    weight: int
+
+    @property
+    def as_pydcs(self) -> PydcsWeapon:
+        return {
+            "clsid": self.cls_id,
+            "name": self.name,
+            "weight": self.weight,
+        }
+
+    @classmethod
+    def from_pydcs(cls, weapon_data: PydcsWeapon) -> Weapon:
+        return cls(
+            cast(str, weapon_data["clsid"]),
+            cast(str, weapon_data["name"]),
+            cast(int, weapon_data["weight"])
+        )
+
+
+@dataclass(frozen=True)
+class Pylon:
+    number: int
+    allowed: Set[Weapon]
+
+    def can_equip(self, weapon: Weapon) -> bool:
+        return weapon in self.allowed
+
+    def equip(self, group: FlyingGroup, weapon: Weapon) -> None:
+        if not self.can_equip(weapon):
+            raise ValueError(f"Pylon {self.number} cannot equip {weapon.name}")
+        group.load_pylon(self.make_pydcs_assignment(weapon), self.number)
+
+    def make_pydcs_assignment(self, weapon: Weapon) -> PydcsWeaponAssignment:
+        return self.number, weapon.as_pydcs
+
+    @classmethod
+    def for_aircraft(cls, aircraft: Type[FlyingType], number: int) -> Pylon:
+        # In pydcs these are all arbitrary inner classes of the aircraft type.
+        # The only way to identify them is by their name.
+        pylons = [v for v in aircraft.__dict__.values() if
+                  inspect.isclass(v) and v.__name__.startswith("Pylon")]
+
+        # And that Pylon class has members with irrelevant names that have
+        # values of (pylon number, allowed weapon).
+        allowed = set()
+        for pylon in pylons:
+            for key, value in pylon.__dict__.items():
+                if key.startswith("__"):
+                    continue
+                pylon_number, weapon = value
+                if pylon_number != number:
+                    continue
+                allowed.add(Weapon.from_pydcs(weapon))
+
+        return cls(number, allowed)
+
+    @classmethod
+    def iter_pylons(cls, aircraft: Type[FlyingType]) -> Iterator[Pylon]:
+        for pylon in sorted(list(aircraft.pylons)):
+            yield cls.for_aircraft(aircraft, pylon)

gen/aircraft.py

@@ -74,6 +74,7 @@ from dcs.unittype import FlyingType, UnitType
 
 from game import db
 from game.data.cap_capabilities_db import GUNFIGHTERS
+from game.data.weapons import Pylon
 from game.factions.faction import Faction
 from game.settings import Settings
 from game.theater.controlpoint import (
@@ -902,22 +903,20 @@ class AircraftConflictGenerator:
         else:
             assert False
 
-    def _setup_custom_payload(self, flight, group:FlyingGroup):
-        if flight.use_custom_loadout:
+    @staticmethod
+    def _setup_custom_payload(flight: Flight, group: FlyingGroup) -> None:
+        if not flight.use_custom_loadout:
+            return
 
-            logging.info("Custom loadout for flight : " + flight.__repr__())
-            for p in group.units:
-                p.pylons.clear()
+        logging.info("Custom loadout for flight : " + flight.__repr__())
+        for p in group.units:
+            p.pylons.clear()
 
-            for key in flight.loadout.keys():
-                if "Pylon" + key in flight.unit_type.__dict__.keys():
-                    print(flight.loadout)
-                    weapon_dict = flight.unit_type.__dict__["Pylon" + key].__dict__
-                    if flight.loadout[key] in weapon_dict.keys():
-                        weapon = weapon_dict[flight.loadout[key]]
-                        group.load_pylon(weapon, int(key))
-                else:
-                    logging.warning("Pylon not found ! => Pylon" + key + " on " + str(flight.unit_type))
+        for pylon_number, weapon in flight.loadout.items():
+            if weapon is None:
+                continue
+            pylon = Pylon.for_aircraft(flight.unit_type, pylon_number)
+            pylon.equip(group, weapon)
 
     def clear_parking_slots(self) -> None:
         for cp in self.game.theater.controlpoints:

gen/flights/flight.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections import defaultdict
 from datetime import timedelta
 from enum import Enum
 from typing import Dict, List, Optional, TYPE_CHECKING, Type
@@ -9,6 +10,7 @@ from dcs.point import MovingPoint, PointAction
 from dcs.unittype import FlyingType
 
 from game import db
+from game.data.weapons import Weapon
 from game.theater.controlpoint import ControlPoint, MissionTarget
 from game.utils import Distance, meters
 
@@ -148,7 +150,7 @@ class Flight:
         self.flight_type = flight_type
         # TODO: Replace with FlightPlan.
         self.targets: List[MissionTarget] = []
-        self.loadout: Dict[str, str] = {}
+        self.loadout: Dict[int, Optional[Weapon]] = {}
         self.start_type = start_type
         self.use_custom_loadout = False
         self.client_count = 0

qt_ui/windows/mission/flight/payload/QLoadoutEditor.py

@@ -2,6 +2,7 @@ import inspect
 
 from PySide2.QtWidgets import QLabel, QHBoxLayout, QGroupBox, QSpinBox, QGridLayout, QVBoxLayout, QSizePolicy
 
+from game.data.weapons import Pylon
 from qt_ui.windows.mission.flight.payload.QPylonEditor import QPylonEditor
 
 
@@ -19,16 +20,12 @@ class QLoadoutEditor(QGroupBox):
         hboxLayout = QVBoxLayout(self)
         layout = QGridLayout(self)
 
-        pylons = [v for v in self.flight.unit_type.__dict__.values() if inspect.isclass(v) and v.__name__.startswith("Pylon")]
-        for i, pylon in enumerate(pylons):
-            label = QLabel("<b>{}</b>".format(pylon.__name__[len("Pylon"):]))
-            label.setSizePolicy(QSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed))
+        for i, pylon in enumerate(Pylon.iter_pylons(self.flight.unit_type)):
+            label = QLabel(f"<b>{pylon.number}</b>")
+            label.setSizePolicy(
+                QSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed))
             layout.addWidget(label, i, 0)
-            try:
-                pylon_number = int(pylon.__name__.split("Pylon")[1])
-            except:
-                pylon_number = i+1
-            layout.addWidget(QPylonEditor(flight, pylon, pylon_number), i, 1)
+            layout.addWidget(QPylonEditor(flight, pylon), i, 1)
 
         hboxLayout.addLayout(layout)
         hboxLayout.addStretch()

qt_ui/windows/mission/flight/payload/QPylonEditor.py

@@ -1,38 +1,37 @@
 import logging
+import operator
+from typing import Optional
 
-from PySide2.QtWidgets import QWidget, QSpinBox, QComboBox
+from PySide2.QtWidgets import QComboBox
+
+from game.data.weapons import Pylon, Weapon
+from gen.flights.flight import Flight
 
 
 class QPylonEditor(QComboBox):
 
-    def __init__(self, flight, pylon, pylon_number):
-        super(QPylonEditor, self).__init__()
-        self.pylon = pylon
-        self.pylon_number = pylon_number
+    def __init__(self, flight: Flight, pylon: Pylon) -> None:
+        super().__init__()
         self.flight = flight
+        self.pylon = pylon
 
-        self.possible_loadout = [i for i in self.pylon.__dict__.keys() if i[:2] != '__']
+        current = self.flight.loadout.get(self.pylon.number)
 
-        if not str(self.pylon_number) in self.flight.loadout.keys():
-            self.flight.loadout[str(self.pylon_number)] = ""
-
-        self.addItem("None")
-        for i,k in enumerate(self.possible_loadout):
-            self.addItem(str(self.pylon.__dict__[k][1]["name"]))
-            if self.flight.loadout[str(self.pylon_number)] == str(k):
+        self.addItem("None", None)
+        allowed = sorted(pylon.allowed, key=operator.attrgetter("name"))
+        for i, weapon in enumerate(allowed):
+            self.addItem(weapon.name, weapon)
+            if current == weapon:
                 self.setCurrentIndex(i + 1)
 
-        self.currentTextChanged.connect(self.on_pylon_change)
+        self.currentIndexChanged.connect(self.on_pylon_change)
 
     def on_pylon_change(self):
-        selected = self.currentText()
-        if selected == "None":
-            logging.info("Pylon " + str(self.pylon_number) + " emptied")
-            self.flight.loadout[str(self.pylon_number)] = ""
-        else:
-            logging.info("Pylon " + str(self.pylon_number) + " changed to " + selected)
-            for i, k in enumerate(self.possible_loadout):
-                if selected == str(self.pylon.__dict__[k][1]["name"]):
-                    self.flight.loadout[str(self.pylon_number)] = str(k)
-                    break
+        selected: Optional[Weapon] = self.currentData()
+        self.flight.loadout[self.pylon.number] = selected
 
+        if selected is None:
+            logging.debug(f"Pylon {self.pylon.number} emptied")
+        else:
+            logging.debug(
+                f"Pylon {self.pylon.number} changed to {selected.name}")